SimplifyGlobals: Optimize away globals that are only read in order to write themselves (#4070)

If the only uses of a global are

if (global == 0) {
  global = 1;
}

Then we do not need that global: while it has both reads and
writes, the value in the global does not cause anything observable.
It is read, but only to write to itself, and nothing else.

This happens in real-world code from j2cl quite a lot, as they have
an initialization pattern with globals, and in some cases we can
optimize away the work done in the initialization, leaving only the
globals in this pattern.

2 files changed, 398 insertions, 14 deletions

diff --git a/src/passes/SimplifyGlobals.cpp b/src/passes/SimplifyGlobals.cpp
index b57490d3c..5a6337bd6 100644
--- a/src/passes/SimplifyGlobals.cpp
+++ b/src/passes/SimplifyGlobals.cpp
@@ -25,6 +25,8 @@
 //  * Apply the constant values of previous global.sets, in a linear
 //    execution trace.
 //  * Remove writes to globals that are never read from.
+//  * Remove writes to globals that are only read from in order to write (see
+//    below, "readOnlyToWrite").
 //
 // Some globals may not have uses after these changes, which we leave
 // to other passes to optimize.
@@ -50,10 +52,36 @@ namespace wasm {
 namespace {
 
 struct GlobalInfo {
+  // Whether the global is imported and exported.
   bool imported = false;
   bool exported = false;
-  std::atomic<bool> written;
-  std::atomic<bool> read;
+
+  // How many times the global is written and read.
+  std::atomic<Index> written{0};
+  std::atomic<Index> read{0};
+
+  // How many times the global is "read, but only to write", that is, is used in
+  // this pattern:
+  //
+  //   if (global == X) { global = Y }
+  //
+  // where X and Y have no side effects. If all we have are such reads only to
+  // write then the global is really not necessary, even though there are both
+  // reads and writes of it.
+  //
+  // This pattern can show up in global initialization code, where in the block
+  // alongside "global = Y" there was some useful code, but the optimizer
+  // managed to remove it. For example,
+  //
+  //   if (global == 0) { global = 1; sideEffect(); }
+  //
+  // If the global's initial value is the default 0, and there are no other uses
+  // of this global, then this code will run sideEffect() the very first time we
+  // reach here. We therefore need to keep this global and its reads and writes.
+  // However, if sideEffect() were removed, then we read the global only to
+  // write it - and nothing else - and so we can optimize away that global
+  // entirely.
+  std::atomic<Index> readOnlyToWrite{0};
 };
 
 using GlobalInfoMap = std::map<Name, GlobalInfo>;
@@ -65,9 +93,53 @@ struct GlobalUseScanner : public WalkerPass<PostWalker<GlobalUseScanner>> {
 
   GlobalUseScanner* create() override { return new GlobalUseScanner(infos); }
 
-  void visitGlobalSet(GlobalSet* curr) { (*infos)[curr->name].written = true; }
+  void visitGlobalSet(GlobalSet* curr) { (*infos)[curr->name].written++; }
+
+  void visitGlobalGet(GlobalGet* curr) { (*infos)[curr->name].read++; }
+
+  void visitIf(If* curr) {
+    // We are looking for
+    //
+    //   if (global == X) { global = Y }
+    //
+    // Ignore an if-else, which cannot be that.
+    if (curr->ifFalse) {
+      return;
+    }
+
+    // See if reading a specific global is the only effect the condition has.
+    EffectAnalyzer condition(
+      getPassOptions(), getModule()->features, curr->condition);
+
+    if (condition.globalsRead.size() != 1) {
+      return;
+    }
+    auto global = *condition.globalsRead.begin();
+    condition.globalsRead.clear();
+    if (condition.hasAnything()) {
+      return;
+    }
+
+    // See if writing the same global is the only effect the body has. (Note
+    // that we don't need to care about the case where the body has no effects
+    // at all - other pass would handle that trivial situation.)
+    EffectAnalyzer ifTrue(
+      getPassOptions(), getModule()->features, curr->ifTrue);
+    if (ifTrue.globalsWritten.size() != 1) {
+      return;
+    }
+    auto writtenGlobal = *ifTrue.globalsWritten.begin();
+    if (writtenGlobal != global) {
+      return;
+    }
+    ifTrue.globalsWritten.clear();
+    if (ifTrue.hasAnything()) {
+      return;
+    }
 
-  void visitGlobalGet(GlobalGet* curr) { (*infos)[curr->name].read = true; }
+    // This is exactly the pattern we sought!
+    (*infos)[global].readOnlyToWrite++;
+  }
 
 private:
   GlobalInfoMap* infos;
@@ -215,18 +287,29 @@ struct SimplifyGlobals : public Pass {
     runner = runner_;
     module = module_;
 
+    while (iteration()) {
+    }
+  }
+
+  bool iteration() {
     analyze();
 
-    removeWritesToUnreadGlobals();
+    // Removing unneeded writes can in some cases lead to more optimizations
+    // that we need an entire additional iteration to perform, see below.
+    bool more = removeUnneededWrites();
 
     preferEarlierImports();
 
     propagateConstantsToGlobals();
 
     propagateConstantsToCode();
+
+    return more;
   }
 
   void analyze() {
+    map.clear();
+
     // First, find out all the relevant info.
     for (auto& global : module->globals) {
       auto& info = map[global->name];
@@ -250,21 +333,82 @@ struct SimplifyGlobals : public Pass {
     }
   }
 
-  void removeWritesToUnreadGlobals() {
-    // Globals that are not exports and not read from can be eliminated
-    // (even if they are written to).
-    NameSet unreadGlobals;
+  // Removes writes from globals that will never do anything useful with the
+  // written value anyhow. Returns whether an addition iteration is necessary.
+  bool removeUnneededWrites() {
+    bool more = false;
+
+    // Globals that are not exports and not read from are unnecessary (even if
+    // they are written to). Likewise, globals that are only read from in order
+    // to write to themselves are unnecessary. First, find such globals.
+    NameSet unnecessaryGlobals;
     for (auto& global : module->globals) {
       auto& info = map[global->name];
-      if (!info.imported && !info.exported && !info.read) {
-        unreadGlobals.insert(global->name);
+
+      if (!info.written) {
+        // No writes occur here, so there is nothing for us to remove.
+        continue;
+      }
+
+      if (info.imported || info.exported) {
+        // If the global is observable from the outside, we can't do anythng
+        // here.
+        //
+        // TODO: optimize the case of an imported but immutable global, etc.
+        continue;
+      }
+
+      // We only ever optimize read-only-to-write if all of our reads are done
+      // in places we identified as read-only-to-write. That is, we have
+      // eliminated the possibility of any other uses. (Technically, each
+      // read-to-write location might have more than one read since we did not
+      // count them, but only verified there was one read or more; but this is
+      // good enough as the common case has exactly one.)
+      //
+      // Note that there might be more writes, if there are additional writes
+      // besides those in the read-only-to-write locations. But we can ignore
+      // those, as whatever they write will not be read in order to do anything
+      // of value.
+      bool onlyReadOnlyToWrite = (info.read == info.readOnlyToWrite);
+
+      // There is at least one write in each read-only-to-write location, unless
+      // our logic is wrong somewhere.
+      assert(info.written >= info.readOnlyToWrite);
+
+      if (!info.read || onlyReadOnlyToWrite) {
+        unnecessaryGlobals.insert(global->name);
+
         // We can now mark this global as immutable, and un-written, since we
-        // are about to remove all the `.set` operations on it.
+        // are about to remove all the operations on it.
         global->mutable_ = false;
-        info.written = false;
+        info.written = 0;
+
+        // Nested old-read-to-write expressions require another full iteration
+        // to optimize, as we have:
+        //
+        //   if (a) {
+        //     a = 1;
+        //     if (b) {
+        //       b = 1;
+        //     }
+        //   }
+        //
+        // The first iteration can only optimize b, as the outer if's body has
+        // more effects than we understand. After finishing the first iteration,
+        // b will no longer exist, removing those effects.
+        if (onlyReadOnlyToWrite) {
+          more = true;
+        }
       }
     }
-    GlobalSetRemover(&unreadGlobals, optimize).run(runner, module);
+
+    // Remove all the sets on the unnecessary globals. Later optimizations can
+    // then see that since the global has no writes, it is a constant, which
+    // will lead to removal of gets, and after removing them, the global itself
+    // will be removed as well.
+    GlobalSetRemover(&unnecessaryGlobals, optimize).run(runner, module);
+
+    return more;
   }
 
   void preferEarlierImports() {
diff --git a/test/lit/passes/simplify-globals-read_only_to_write.wast b/test/lit/passes/simplify-globals-read_only_to_write.wast
new file mode 100644
index 000000000..1ba0b9ecc
--- /dev/null
+++ b/test/lit/passes/simplify-globals-read_only_to_write.wast
@@ -0,0 +1,240 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; NOTE: This test was ported using port_test.py and could be cleaned up.
+
+;; RUN: foreach %s %t wasm-opt --simplify-globals -S -o - | filecheck %s
+
+;; A global that is only read in order to be written is not needed.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $global i32 (i32.const 0))
+  (global $global (mut i32) (i32.const 0))
+  ;; CHECK:      (func $simple
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $simple
+    (if
+      (global.get $global)
+      (global.set $global (i32.const 1))
+    )
+  )
+  ;; CHECK:      (func $more-with-no-side-effects
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.eqz
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (block $block
+  ;; CHECK-NEXT:    (nop)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $more-with-no-side-effects
+    (if
+      ;; Also test for other operations in the condition, with no effects.
+      (i32.eqz
+        (global.get $global)
+      )
+      ;; Also test for other operations in the body, with no effects.
+      (block
+        (nop)
+        (global.set $global (i32.const 1))
+      )
+    )
+  )
+  ;; CHECK:      (func $additional-set
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $additional-set
+    ;; An additional set does not prevent this optimization: the value written
+    ;; will never be read in a way that matters.
+    (global.set $global (i32.const 2))
+  )
+)
+;; An additional read prevents the read-only-to-write optimization.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $global (mut i32) (i32.const 0))
+  (global $global (mut i32) (i32.const 0))
+  ;; CHECK:      (func $additional-read
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (global.get $global)
+  ;; CHECK-NEXT:   (global.set $global
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (global.get $global)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $additional-read
+    (if
+      (global.get $global)
+      (global.set $global (i32.const 1))
+    )
+    (drop
+      (global.get $global)
+    )
+  )
+)
+;; We do not optimize if-elses in the read-only-to-write optimization.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $global (mut i32) (i32.const 0))
+  (global $global (mut i32) (i32.const 0))
+  ;; CHECK:      (func $if-else
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (global.get $global)
+  ;; CHECK-NEXT:   (global.set $global
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (nop)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $if-else
+    (if
+      (global.get $global)
+      (global.set $global (i32.const 1))
+      (nop)
+    )
+  )
+)
+;; Side effects in the condition prevent the read-only-to-write optimization.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $global (mut i32) (i32.const 0))
+  (global $global (mut i32) (i32.const 0))
+  ;; CHECK:      (global $other (mut i32) (i32.const 0))
+  (global $other (mut i32) (i32.const 0))
+  ;; CHECK:      (func $side-effects-in-condition
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (block $block (result i32)
+  ;; CHECK-NEXT:    (global.set $other
+  ;; CHECK-NEXT:     (i32.const 2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (global.get $global)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (global.set $global
+  ;; CHECK-NEXT:    (i32.const 1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $side-effects-in-condition
+    (if
+      (block (result i32)
+        (global.set $other (i32.const 2))
+        (drop (global.get $other))
+        (global.get $global)
+      )
+      (global.set $global (i32.const 1))
+    )
+  )
+)
+;; Side effects in the body prevent the read-only-to-write optimization.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $global (mut i32) (i32.const 0))
+  (global $global (mut i32) (i32.const 0))
+  ;; CHECK:      (global $other (mut i32) (i32.const 0))
+  (global $other (mut i32) (i32.const 0))
+  ;; CHECK:      (func $side-effects-in-body
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (global.get $global)
+  ;; CHECK-NEXT:   (block $block
+  ;; CHECK-NEXT:    (global.set $global
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (global.set $other
+  ;; CHECK-NEXT:     (i32.const 2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 2)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $side-effects-in-body
+    (if
+      (global.get $global)
+      (block
+        (global.set $global (i32.const 1))
+        (global.set $other (i32.const 2))
+        (drop (global.get $other))
+      )
+    )
+  )
+)
+;; Nested patterns work as well, in a single run of the pass.
+(module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (global $a i32 (i32.const 0))
+  (global $a (mut i32) (i32.const 0))
+  ;; CHECK:      (global $b i32 (i32.const 0))
+  (global $b (mut i32) (i32.const 0))
+  ;; CHECK:      (global $c i32 (i32.const 0))
+  (global $c (mut i32) (i32.const 0))
+
+  ;; CHECK:      (func $nested
+  ;; CHECK-NEXT:  (if
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:   (block $block
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (i32.const 1)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (if
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:     (block $block1
+  ;; CHECK-NEXT:      (if
+  ;; CHECK-NEXT:       (i32.const 0)
+  ;; CHECK-NEXT:       (block $block3
+  ;; CHECK-NEXT:        (drop
+  ;; CHECK-NEXT:         (i32.const 2)
+  ;; CHECK-NEXT:        )
+  ;; CHECK-NEXT:       )
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:      (drop
+  ;; CHECK-NEXT:       (i32.const 3)
+  ;; CHECK-NEXT:      )
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $nested
+    (if
+      (global.get $a)
+      (block
+        (global.set $a (i32.const 1))
+        (if
+          (global.get $b)
+          (block
+            (if
+              (global.get $c)
+              (block
+                (global.set $c (i32.const 2))
+              )
+            )
+            (global.set $b (i32.const 3))
+          )
+        )
+      )
+    )
+  )
+)