Asyncify: Instrument indirect calls from functions in add-list or onl…

…y-list (#2913)

When doing manual tuning of calls using asyncify lists, we want it to
be possible to write out all the functions that can be on the stack when
pausing, and for that to work. This did not quite work right with the
ignore-indirect option: that would ignore all indirect calls all the
time, so that if foo() calls bar() indirectly, that indirect call was
not instrumented (we didn't check for a pause around it), even if
both foo() and bar() were listed. There was no way to make that
work (except for not ignoring indirect calls at all).

This PR makes the add-list and only-lists fully instrument the functions
mentioned in them: both themselves, and indirect calls from them.
(Note that direct calls need no special handling - we can just add
the direct call target to the add-list or only-list.)

This may add some overhead to existing users, but only in a function
that is instrumented anyhow, and also indirect calls are slow anyhow,
so it's probably fine. And it is simpler to do it this way instead of
adding another list for indirect call handling.

src/passes/Asyncify.cpp

@@ -265,6 +265,27 @@
 //      If the "only-list" is provided, then *only* the functions in the list
 //      will be instrumented, and nothing else.
 //
+// Note that there are two types of instrumentation that happen for each
+// function: if foo() can be part of a pause/resume operation, then we need to
+// instrument code inside it to support pausing and resuming, but also, we need
+// callers of the function to instrument calls to it. Normally this is already
+// taken care of as the callers need to be instrumented as well anyhow. However,
+// the ignore-indirect option makes things more complicated, since we can't tell
+// where all the calls to foo() are - all we see are indirect calls that do not
+// refer to foo() by name. To make it possible for you to use the add-list or
+// only-list with ignore-indirect, those lists affect *both* kinds of
+// instrumentation. That is, if parent() calls foo() indirectly, and you add
+// parent() to the add-list, then the indirect calls in parent() will be
+// instrumented to support pausing/resuming, even if ignore-indirect is set.
+// Typically you don't need to think about this, and just add all the functions
+// that can be on the stack while pausing - what this means is that when you do
+// so, indirect calls will just work. (The cost is that an instrumented function
+// will check for pausing at all indirect calls, which may be unnecessary in
+// some cases; but this is an instrumented function anyhow, and indirect calls
+// are slower anyhow, so this simple model seems good enough - a more complex
+// model where you can specify "instrument, but not indirect calls from me"
+// would likely have little benefit.)
+//
 // TODO When wasm has GC, extending the live ranges of locals can keep things
 //      alive unnecessarily. We may want to set locals to null at the end
 //      of their original range.
@@ -478,6 +499,7 @@ class ModuleAnalyzer {
     // that call it are instrumented. This is not done for the bottom.
     bool isTopMostRuntime = false;
     bool inRemoveList = false;
+    bool addedFromList = false;
   };
 
   typedef std::map<Function*, Info> Map;
@@ -634,15 +656,22 @@ class ModuleAnalyzer {
       // Only the functions in the only-list can change the state.
       for (auto& func : module.functions) {
         if (!func->imported()) {
-          map[func.get()].canChangeState = onlyList.match(func->name);
+          auto& info = map[func.get()];
+          bool matched = onlyList.match(func->name);
+          info.canChangeState = matched;
+          if (matched) {
+            info.addedFromList = true;
+          }
         }
       }
     }
 
     if (!addListInput.empty()) {
       for (auto& func : module.functions) {
         if (!func->imported() && addList.match(func->name)) {
-          map[func.get()].canChangeState = true;
+          auto& info = map[func.get()];
+          info.canChangeState = true;
+          info.addedFromList = true;
         }
       }
     }
@@ -657,7 +686,7 @@ class ModuleAnalyzer {
     return info.canChangeState && !info.isTopMostRuntime;
   }
 
-  bool canChangeState(Expression* curr) {
+  bool canChangeState(Expression* curr, Function* func) {
     // Look inside to see if we call any of the things we know can change the
     // state.
     // TODO: caching, this is O(N^2)
@@ -683,29 +712,30 @@ class ModuleAnalyzer {
           canChangeState = true;
         }
       }
-      void visitCallIndirect(CallIndirect* curr) {
-        if (canIndirectChangeState) {
-          canChangeState = true;
-        }
-        // TODO optimize the other case, at least by type
-      }
+      void visitCallIndirect(CallIndirect* curr) { hasIndirectCall = true; }
       Module* module;
       ModuleAnalyzer* analyzer;
       Map* map;
-      bool canIndirectChangeState;
+      bool hasIndirectCall = false;
       bool canChangeState = false;
       bool isBottomMostRuntime = false;
     };
     Walker walker;
     walker.module = &module;
     walker.analyzer = this;
     walker.map = &map;
-    walker.canIndirectChangeState = canIndirectChangeState;
     walker.walk(curr);
-    if (walker.isBottomMostRuntime) {
-      walker.canChangeState = false;
+    // An indirect call is normally ignored if we are ignoring indirect calls.
+    // However, see the docs at the top: if the function we are inside was
+    // specifically added by the user (in the only-list or the add-list) then we
+    // instrument indirect calls from it (this allows specifically allowing some
+    // indirect calls but not others).
+    if (walker.hasIndirectCall &&
+        (canIndirectChangeState || map[func].addedFromList)) {
+      walker.canChangeState = true;
     }
-    return walker.canChangeState;
+    // The bottom-most runtime can never change the state.
+    return walker.canChangeState && !walker.isBottomMostRuntime;
   }
 
   FakeGlobalHelper fakeGlobals;
@@ -814,7 +844,7 @@ struct AsyncifyFlow : public Pass {
   Index callIndex = 0;
 
   Expression* process(Expression* curr) {
-    if (!analyzer->canChangeState(curr)) {
+    if (!analyzer->canChangeState(curr, func)) {
       return makeMaybeSkip(curr);
     }
     // The IR is in flat form, which makes this much simpler: there are no
@@ -856,12 +886,13 @@ struct AsyncifyFlow : public Pass {
       Index i = 0;
       auto& list = block->list;
       while (i < list.size()) {
-        if (analyzer->canChangeState(list[i])) {
+        if (analyzer->canChangeState(list[i], func)) {
           list[i] = process(list[i]);
           i++;
         } else {
           Index end = i + 1;
-          while (end < list.size() && !analyzer->canChangeState(list[end])) {
+          while (end < list.size() &&
+                 !analyzer->canChangeState(list[end], func)) {
             end++;
           }
           // We have a range of [i, end) in which the state cannot change,
@@ -886,7 +917,7 @@ struct AsyncifyFlow : public Pass {
     } else if (auto* iff = curr->dynCast<If>()) {
       // The state change cannot be in the condition due to flat form, so it
       // must be in one of the children.
-      assert(!analyzer->canChangeState(iff->condition));
+      assert(!analyzer->canChangeState(iff->condition, func));
       // We must linearize this, which means we pass through both arms if we
       // are rewinding.
       if (!iff->ifFalse) {

test/passes/asyncify_pass-arg=asyncify-addlist@foo_pass-arg=asyncify-ignore-indirect.txt

@@ -0,0 +1,195 @@
+(module
+ (type $none_=>_none (func))
+ (type $i32_=>_none (func (param i32)))
+ (type $none_=>_i32 (func (result i32)))
+ (import "env" "import" (func $import))
+ (memory $0 1 2)
+ (table $0 1 funcref)
+ (global $__asyncify_state (mut i32) (i32.const 0))
+ (global $__asyncify_data (mut i32) (i32.const 0))
+ (export "asyncify_start_unwind" (func $asyncify_start_unwind))
+ (export "asyncify_stop_unwind" (func $asyncify_stop_unwind))
+ (export "asyncify_start_rewind" (func $asyncify_start_rewind))
+ (export "asyncify_stop_rewind" (func $asyncify_stop_rewind))
+ (export "asyncify_get_state" (func $asyncify_get_state))
+ (func $foo
+  (local $0 i32)
+  (local $1 i32)
+  (if
+   (i32.eq
+    (global.get $__asyncify_state)
+    (i32.const 2)
+   )
+   (nop)
+  )
+  (local.set $0
+   (block $__asyncify_unwind (result i32)
+    (block
+     (block
+      (if
+       (i32.eq
+        (global.get $__asyncify_state)
+        (i32.const 2)
+       )
+       (block
+        (i32.store
+         (global.get $__asyncify_data)
+         (i32.add
+          (i32.load
+           (global.get $__asyncify_data)
+          )
+          (i32.const -4)
+         )
+        )
+        (local.set $1
+         (i32.load
+          (i32.load
+           (global.get $__asyncify_data)
+          )
+         )
+        )
+       )
+      )
+      (block
+       (if
+        (i32.eq
+         (global.get $__asyncify_state)
+         (i32.const 0)
+        )
+        (call $nothing)
+       )
+       (if
+        (if (result i32)
+         (i32.eq
+          (global.get $__asyncify_state)
+          (i32.const 0)
+         )
+         (i32.const 1)
+         (i32.eq
+          (local.get $1)
+          (i32.const 0)
+         )
+        )
+        (block
+         (call_indirect (type $none_=>_none)
+          (i32.const 0)
+         )
+         (if
+          (i32.eq
+           (global.get $__asyncify_state)
+           (i32.const 1)
+          )
+          (br $__asyncify_unwind
+           (i32.const 0)
+          )
+         )
+        )
+       )
+      )
+     )
+     (return)
+    )
+   )
+  )
+  (block
+   (i32.store
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (local.get $0)
+   )
+   (i32.store
+    (global.get $__asyncify_data)
+    (i32.add
+     (i32.load
+      (global.get $__asyncify_data)
+     )
+     (i32.const 4)
+    )
+   )
+  )
+  (nop)
+ )
+ (func $bar
+  (call $nothing)
+  (call_indirect (type $none_=>_none)
+   (i32.const 0)
+  )
+ )
+ (func $nothing
+  (nop)
+ )
+ (func $asyncify_start_unwind (param $0 i32)
+  (global.set $__asyncify_state
+   (i32.const 1)
+  )
+  (global.set $__asyncify_data
+   (local.get $0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_stop_unwind
+  (global.set $__asyncify_state
+   (i32.const 0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_start_rewind (param $0 i32)
+  (global.set $__asyncify_state
+   (i32.const 2)
+  )
+  (global.set $__asyncify_data
+   (local.get $0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_stop_rewind
+  (global.set $__asyncify_state
+   (i32.const 0)
+  )
+  (if
+   (i32.gt_u
+    (i32.load
+     (global.get $__asyncify_data)
+    )
+    (i32.load offset=4
+     (global.get $__asyncify_data)
+    )
+   )
+   (unreachable)
+  )
+ )
+ (func $asyncify_get_state (result i32)
+  (global.get $__asyncify_state)
+ )
+)

test/passes/asyncify_pass-arg=asyncify-addlist@foo_pass-arg=asyncify-ignore-indirect.wast

@@ -0,0 +1,18 @@
+(module
+  (type $t (func))
+  (memory 1 2)
+  (table 1 funcref)
+  (elem (i32.const 0))
+  (import "env" "import" (func $import))
+  (func $foo ;; doesn't look like it needs instrumentation, but in add list
+    (call $nothing)
+    (call_indirect (type $t) (i32.const 0))
+  )
+  (func $bar ;; doesn't look like it needs instrumentation, and not in add list
+    (call $nothing)
+    (call_indirect (type $t) (i32.const 0))
+  )
+  (func $nothing
+  )
+)
+