OnceReduction: Optimize idempotent functions

scripts/test/fuzzing.py

@@ -119,6 +119,7 @@
     'idempotent.wast',
     'optimize-instructions_idempotent.wast',
     'duplicate-function-elimination_annotations.wast',
+    'once-reduction_idempotent.wast',
     # Not fully implemented.
     'waitqueue.wast',
 ]

src/passes/OnceReduction.cpp

@@ -40,10 +40,21 @@
 // TODO: "Once" globals are effectively boolean in that all non-zero values are
 //       indistinguishable, and so we could rewrite them all to be 1.
 //
+// In addition, functions marked with @binaryen.idempotent are also "once", just
+// without an explicit global that guards them. That is, when we find such a
+// global, as mentioned above, we effectively prove that they are idempotent,
+// but can also just use the annotation when we find it.
+//
+// TODO: We could mark "once" functions with a global as idempotent, but that
+//       then requires the user to strip the annotations later, so we do not do
+//       this for now.
+//
 
 #include <atomic>
 
 #include "cfg/domtree.h"
+#include "ir/intrinsics.h"
+#include "ir/names.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/unique_deferring_queue.h"
@@ -70,7 +81,14 @@ struct OptInfo {
 
   // Maps functions to whether they are "once", by indicating the global that
   // they use for that purpose. An empty name means they are not "once".
-  std::unordered_map<Name, Name> onceFuncs;
+  //
+  // When we see an idempotent-marked function, which as mentioned above is
+  // effectively "once" but does not have an explicit global controlling it, we
+  // create a fake global name for it to use here. That gives each function we
+  // can optimize a unique global name for our optimizer to track. That is, the
+  // global name is a real global for ones we found a global for, and for
+  // idempotent functions, it is a unique name that is not an actual global.
+  std::unordered_map<Name, Name> onceFuncGlobals;
 
   // For each function, the "once" globals that are definitely set after calling
   // it. If the function is "once" itself, that is included, but it also
@@ -126,7 +144,7 @@ struct Scanner : public WalkerPass<PostWalker<Scanner>> {
         // This is a "once" function, as best we can tell for now. Further
         // information may cause a problem, say, if the global is used in a bad
         // way in another function, so we may undo this.
-        optInfo.onceFuncs.at(curr->name) = global;
+        optInfo.onceFuncGlobals.at(curr->name) = global;
 
         // We can ignore the get in the "once" pattern at the top of the
         // function.
@@ -314,10 +332,10 @@ struct Optimizer
           }
         } else if (auto* call = expr->dynCast<Call>()) {
           auto target = call->target;
-          if (optInfo.onceFuncs.at(target).is()) {
+          if (optInfo.onceFuncGlobals.at(target).is()) {
             // The global used by the "once" func is written.
             assert(call->operands.empty());
-            optimizeOnce(optInfo.onceFuncs.at(target));
+            optimizeOnce(optInfo.onceFuncGlobals.at(target));
             continue;
           }
 
@@ -365,7 +383,7 @@ struct OnceReduction : public Pass {
     }
     for (auto& func : module->functions) {
       // Fill in the map so that it can be operated on in parallel.
-      optInfo.onceFuncs[func->name] = Name();
+      optInfo.onceFuncGlobals[func->name] = Name();
     }
     for (auto& ex : module->exports) {
       if (ex->kind == ExternalKind::Global) {
@@ -382,12 +400,43 @@ struct OnceReduction : public Pass {
     // Combine the information. We found which globals appear to be "once", but
     // other information may have proven they are not so, in fact. Specifically,
     // for a function to be "once" we need its global to also be such.
-    for (auto& [_, onceGlobal] : optInfo.onceFuncs) {
+    for (auto& [_, onceGlobal] : optInfo.onceFuncGlobals) {
       if (onceGlobal.is() && !optInfo.onceGlobals[onceGlobal]) {
         onceGlobal = Name();
       }
     }
 
+    // Use idempotency. If a function is marked idempotent, we can give it a
+    // fake global id so that we can optimize it.
+    for (auto& func : module->functions) {
+      if (func->getNumParams()) {
+        // We do not yet handle parameters. We would need to make sure they are
+        // equal, to optimize, so we'd need to track more than which functions
+        // we've already seen, in the analysis above. TODO
+        continue;
+      }
+      if (func->getResults().size()) {
+        // We do not yet handle results. We can't just nop such a caller, and
+        // instead should save the result from earlier, and reuse it. TODO
+        continue;
+      }
+
+      auto& globalForFunc = optInfo.onceFuncGlobals[func->name];
+      if (globalForFunc) {
+        // This is already known to be "once", and we know a global for it, so
+        // we can do no better.
+        continue;
+      }
+
+      if (Intrinsics::getAnnotations(func.get()).idempotent) {
+        // Pick a name for the fake global to track this function, and write it
+        // into onceFuncGlobals.
+        globalForFunc = Names::getValidGlobalName(*module, func->name);
+        // Also write into onceGlobals, to mark the (fake) global as once.
+        optInfo.onceGlobals[globalForFunc] = true;
+      }
+    }
+
     // Optimize using what we found. Keep iterating while we find things to
     // optimize, which we estimate using a counter of the total number of once
     // globals set by functions: as that increases, it means we are propagating
@@ -403,7 +452,7 @@ struct OnceReduction : public Pass {
       // Either way, at least fill the data structure for parallel operation.
       auto& set = optInfo.onceGlobalsSetInFuncs[func->name];
 
-      auto global = optInfo.onceFuncs[func->name];
+      auto global = optInfo.onceFuncGlobals[func->name];
       if (global.is()) {
         set.insert(global);
         foundOnce = true;
@@ -455,8 +504,11 @@ struct OnceReduction : public Pass {
     // Iterate deterministically on functions, as the order matters (since we
     // make decisions based on previous actions; see below).
     for (auto& func : module->functions) {
-      if (!optInfo.onceFuncs.at(func->name).is()) {
-        // This is not a "once" function.
+      auto global = optInfo.onceFuncGlobals.at(func->name);
+      if (!global || !module->getGlobalOrNull(global)) {
+        // This is not a "once" function, or it is but it has no corresponding
+        // global (which means this is a fake global name, and this is "once"
+        // only because of idempotency).
         continue;
       }
 
@@ -492,7 +544,7 @@ struct OnceReduction : public Pass {
       }
       auto* payload = list[2];
       if (auto* call = payload->dynCast<Call>()) {
-        if (optInfo.onceFuncs.at(call->target).is()) {
+        if (optInfo.onceFuncGlobals.at(call->target).is()) {
           // All this "once" function does is call another. We do not need the
           // early-exit logic in this one, then, because of the following
           // reasoning. We are comparing these forms:

src/wasm.h

@@ -2262,6 +2262,14 @@ struct CodeAnnotation {
   // calls with the same parameters can be assumed to have no effects. If a
   // value is returned, it will be the same value as returned earlier (for the
   // same parameters).
+  //
+  // Note that this differs from related concepts in C,
+  // https://en.cppreference.com/w/c/language/attributes/reproducible.html#Idempotent
+  // There, idempotency is considered compared to the state of the program,
+  // which means that two idempotent calls with some effect in between cannot be
+  // optimized. Here, we do optimize such situations - the only state we care
+  // about is what is passed in via parameters. This allows us to better
+  // optimize things like Java class constructors.
   bool idempotent = false;
 
   bool operator==(const CodeAnnotation& other) const {

test/lit/passes/once-reduction_idempotent.wast

@@ -0,0 +1,168 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+;; RUN: foreach %s %t wasm-opt --once-reduction -all -S -o - | filecheck %s
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (import "a" "b" (func $import (type $0)))
+  (import "a" "b" (func $import))
+
+  ;; CHECK:      (@binaryen.idempotent)
+  ;; CHECK-NEXT: (func $idempotent (type $0)
+  ;; CHECK-NEXT:  (call $import)
+  ;; CHECK-NEXT: )
+  (@binaryen.idempotent)
+  (func $idempotent
+    ;; This function has side effects, but is marked idempotent.
+    (call $import)
+  )
+
+  ;; CHECK:      (func $potent (type $0)
+  ;; CHECK-NEXT:  (call $import)
+  ;; CHECK-NEXT: )
+  (func $potent
+    ;; As above, but not marked as idempotent.
+    (call $import)
+  )
+
+  ;; CHECK:      (func $basics (type $0)
+  ;; CHECK-NEXT:  (call $idempotent)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (call $potent)
+  ;; CHECK-NEXT:  (call $potent)
+  ;; CHECK-NEXT: )
+  (func $basics
+    ;; We can optimize here, using idempotency. The second call goes away.
+    (call $idempotent)
+    (call $idempotent)
+
+    ;; We can't optimize here.
+    (call $potent)
+    (call $potent)
+  )
+
+  ;; CHECK:      (@binaryen.idempotent)
+  ;; CHECK-NEXT: (func $idempotent2 (type $0)
+  ;; CHECK-NEXT:  (call $import)
+  ;; CHECK-NEXT: )
+  (@binaryen.idempotent)
+  (func $idempotent2
+    ;; A second idempotent function.
+    (call $import)
+  )
+
+  ;; CHECK:      (func $mix (type $0)
+  ;; CHECK-NEXT:  (call $idempotent)
+  ;; CHECK-NEXT:  (call $idempotent2)
+  ;; CHECK-NEXT: )
+  (func $mix
+    ;; We can't optimize here: these are idempotent, but different.
+    (call $idempotent)
+    (call $idempotent2)
+  )
+
+  ;; CHECK:      (func $mix2 (type $0)
+  ;; CHECK-NEXT:  (call $idempotent2)
+  ;; CHECK-NEXT:  (call $idempotent)
+  ;; CHECK-NEXT: )
+  (func $mix2
+    ;; As above, reverse order. We can't optimize here.
+    (call $idempotent2)
+    (call $idempotent)
+  )
+
+  ;; CHECK:      (func $mix3 (type $0)
+  ;; CHECK-NEXT:  (call $idempotent)
+  ;; CHECK-NEXT:  (call $idempotent2)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
+  (func $mix3
+    (call $idempotent)
+    (call $idempotent2)
+
+    ;; We can remove both of these second calls.
+    (call $idempotent)
+    (call $idempotent2)
+
+    ;; These as well, which are reordered.
+    (call $idempotent2)
+    (call $idempotent)
+  )
+)
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (param i32)))
+
+  ;; CHECK:      (import "a" "b" (func $import (type $0)))
+  (import "a" "b" (func $import))
+
+  ;; CHECK:      (@binaryen.idempotent)
+  ;; CHECK-NEXT: (func $idempotent (type $1) (param $x i32)
+  ;; CHECK-NEXT:  (call $import)
+  ;; CHECK-NEXT: )
+  (@binaryen.idempotent)
+  (func $idempotent (param $x i32)
+    ;; As above, but now with a param.
+    (call $import)
+  )
+
+  ;; CHECK:      (func $basics (type $0)
+  ;; CHECK-NEXT:  (call $idempotent
+  ;; CHECK-NEXT:   (i32.const 10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (call $idempotent
+  ;; CHECK-NEXT:   (i32.const 10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $basics
+    ;; We could optimize here in theory, but the pass does not handle params
+    ;; yet. It could compare them for equality, check generativity, etc.
+    (call $idempotent (i32.const 10))
+    (call $idempotent (i32.const 10))
+  )
+)
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (type $1 (func (result i32)))
+
+  ;; CHECK:      (import "a" "b" (func $import (type $0)))
+  (import "a" "b" (func $import))
+
+  ;; CHECK:      (@binaryen.idempotent)
+  ;; CHECK-NEXT: (func $idempotent (type $1) (result i32)
+  ;; CHECK-NEXT:  (call $import)
+  ;; CHECK-NEXT:  (i32.const 42)
+  ;; CHECK-NEXT: )
+  (@binaryen.idempotent)
+  (func $idempotent (result i32)
+    ;; As above, but now with a result.
+    (call $import)
+    (i32.const 42)
+  )
+
+  ;; CHECK:      (func $basics (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $idempotent)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (call $idempotent)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $basics
+    ;; We could optimize here in theory, but the pass does not handle results
+    ;; yet.
+    (drop
+      (call $idempotent)
+    )
+    (drop
+      (call $idempotent)
+    )
+  )
+)