[Wasm GC] Refine function parameter types (#4014)

If a function is always called with a more specific type than it is declared, we can
make the type more specific.

DeadArgumentElimination's name is becoming increasingly misleading, and should
maybe be renamed. But it is the right place for this as it already does an LTO
scan of the call graph and builds up parameter data structures etc.

2 files changed, 310 insertions, 8 deletions

diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index a95d4f91a..a2231469d 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -15,10 +15,8 @@
  */
 
 //
-// Optimizes call arguments in a whole-program manner, removing ones
-// that are not used (dead).
-//
-// Specifically, this does these things:
+// Optimizes call arguments in a whole-program manner. In particular, this
+// removes ones that are not used (dead), but it also does more things:
 //
 //  * Find functions for whom an argument is always passed the same
 //    constant. If so, we can just set that local to that constant
@@ -28,6 +26,8 @@
 //    the previous point was true for an argument, then the second
 //    must as well.)
 //  * Find return values ("return arguments" ;) that are never used.
+//  * Refine the types of arguments, that is make the argument type more
+//    specific if all the passed values allow that.
 //
 // This pass does not depend on flattening, but it may be more effective,
 // as then call arguments never have side effects (which we need to
@@ -40,6 +40,7 @@
 #include "cfg/cfg-traversal.h"
 #include "ir/effects.h"
 #include "ir/element-utils.h"
+#include "ir/find_all.h"
 #include "ir/module-utils.h"
 #include "ir/type-updating.h"
 #include "pass.h"
@@ -307,18 +308,22 @@ struct DAE : public Pass {
         allDroppedCalls[pair.first] = pair.second;
       }
     }
-    // We now have a mapping of all call sites for each function. Check which
-    // are always passed the same constant for a particular argument.
+    // We now have a mapping of all call sites for each function, and can look
+    // for optimization opportunities.
     for (auto& pair : allCalls) {
       auto name = pair.first;
-      // We can only optimize if we see all the calls and can modify
-      // them.
+      // We can only optimize if we see all the calls and can modify them.
       if (infoMap[name].hasUnseenCalls) {
         continue;
       }
       auto& calls = pair.second;
       auto* func = module->getFunction(name);
       auto numParams = func->getNumParams();
+      // Refine argument types before doing anything else. This does not
+      // affect whether an argument is used or not, it just refines the type
+      // where possible.
+      refineArgumentTypes(func, calls, module);
+      // Check if all calls pass the same constant for a particular argument.
       for (Index i = 0; i < numParams; i++) {
         Literal value;
         for (auto* call : calls) {
@@ -515,6 +520,77 @@ private:
       }
     }
   }
+
+  // Given a function and all the calls to it, see if we can refine the type of
+  // its arguments. If we only pass in a subtype, we may as well refine the type
+  // to that.
+  //
+  // This assumes that the function has no calls aside from |calls|, that is, it
+  // is not exported or called from the table or by reference.
+  void refineArgumentTypes(Function* func,
+                           const std::vector<Call*>& calls,
+                           Module* module) {
+    if (!module->features.hasGC()) {
+      return;
+    }
+    auto numParams = func->getNumParams();
+    std::vector<Type> newParamTypes;
+    newParamTypes.reserve(numParams);
+    for (Index i = 0; i < numParams; i++) {
+      auto originalType = func->getLocalType(i);
+      if (!originalType.isRef()) {
+        newParamTypes.push_back(originalType);
+        continue;
+      }
+      Type refinedType = Type::unreachable;
+      for (auto* call : calls) {
+        auto* operand = call->operands[i];
+        refinedType = Type::getLeastUpperBound(refinedType, operand->type);
+        if (refinedType == originalType) {
+          // We failed to refine this parameter to anything more specific.
+          break;
+        }
+      }
+
+      // Nothing is sent here at all; leave such optimizations to DCE.
+      if (refinedType == Type::unreachable) {
+        return;
+      }
+      newParamTypes.push_back(refinedType);
+    }
+
+    // Check if we are able to optimize here before we do the work to scan the
+    // function body.
+    if (Type(newParamTypes) == func->getParams()) {
+      return;
+    }
+
+    // In terms of parameters, we can do this. However, we must also check
+    // local operations in the body, as if the parameter is reused and written
+    // to, then those types must be taken into account as well.
+    for (auto* set : FindAll<LocalSet>(func->body).list) {
+      auto index = set->index;
+      if (func->isParam(index) &&
+          !Type::isSubType(set->value->type, newParamTypes[index])) {
+        // TODO: we could still optimize here, by creating a new local.
+        newParamTypes[index] = func->getLocalType(index);
+      }
+    }
+
+    auto newParams = Type(newParamTypes);
+    if (newParams == func->getParams()) {
+      return;
+    }
+
+    // We can do this! Update the types, including the types of gets.
+    func->setParams(newParams);
+    for (auto* get : FindAll<LocalGet>(func->body).list) {
+      auto index = get->index;
+      if (func->isParam(index)) {
+        get->type = func->getLocalType(index);
+      }
+    }
+  }
 };
 
 Pass* createDAEPass() { return new DAE(); }
diff --git a/test/lit/passes/dae-gc.wast b/test/lit/passes/dae-gc.wast
index 65d65d4c9..b3f8eddd2 100644
--- a/test/lit/passes/dae-gc.wast
+++ b/test/lit/passes/dae-gc.wast
@@ -2,8 +2,18 @@
 ;; RUN: wasm-opt %s -all --dae -S -o - | filecheck %s
 
 (module
+ ;; CHECK:      (type ${i32} (struct (field i32)))
+
  ;; CHECK:      (type ${} (struct ))
  (type ${} (struct))
+ (type ${i32} (struct (field i32)))
+ ;; CHECK:      (type ${i32_i64} (struct (field i32) (field i64)))
+
+ ;; CHECK:      (type ${f64} (struct (field f64)))
+ (type ${f64} (struct (field f64)))
+ (type ${i32_i64} (struct (field i32) (field i64)))
+ ;; CHECK:      (type ${i32_f32} (struct (field i32) (field f32)))
+ (type ${i32_f32} (struct (field i32) (field f32)))
 
  ;; CHECK:      (func $foo
  ;; CHECK-NEXT:  (call $bar)
@@ -67,4 +77,220 @@
    (rtt.canon ${})
   )
  )
+
+ ;; CHECK:      (func $call-various-params-no
+ ;; CHECK-NEXT:  (call $various-params-no
+ ;; CHECK-NEXT:   (ref.null ${})
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $various-params-no
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:   (ref.null ${f64})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-no
+  ;; The first argument gets {} and {i32}; the second {i32} and {f64}; none of
+  ;; those pairs can be optimized.
+  (call $various-params-no
+   (ref.null ${})
+   (ref.null ${i32})
+  )
+  (call $various-params-no
+   (ref.null ${i32})
+   (ref.null ${f64})
+  )
+ )
+ ;; This function is called in ways that do not allow us to alter the types of
+ ;; its parameters (see last function).
+ ;; CHECK:      (func $various-params-no (param $x (ref null ${})) (param $y (ref null ${}))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-no (param $x (ref null ${})) (param $y (ref null ${}))
+  ;; "Use" the locals to avoid other optimizations kicking in.
+  (drop (local.get $x))
+  (drop (local.get $y))
+ )
+
+ ;; CHECK:      (func $call-various-params-yes
+ ;; CHECK-NEXT:  (call $various-params-yes
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $various-params-yes
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:   (i32.const 1)
+ ;; CHECK-NEXT:   (ref.null ${i32_i64})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-yes
+  ;; The first argument gets {i32} and {i32}; the second {i32} and {i32_i64};
+  ;; both of those pairs can be optimized to {i32}.
+  ;; There is also an i32 in the middle, which should not confuse us.
+  (call $various-params-yes
+   (ref.null ${i32})
+   (i32.const 0)
+   (ref.null ${i32})
+  )
+  (call $various-params-yes
+   (ref.null ${i32})
+   (i32.const 1)
+   (ref.null ${i32_i64})
+  )
+ )
+ ;; This function is called in ways that *do* allow us to alter the types of
+ ;; its parameters (see last function).
+ ;; CHECK:      (func $various-params-yes (param $x (ref null ${i32})) (param $i i32) (param $y (ref null ${i32}))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $i)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-yes (param $x (ref null ${})) (param $i i32) (param $y (ref null ${}))
+  ;; "Use" the locals to avoid other optimizations kicking in.
+  (drop (local.get $x))
+  (drop (local.get $i))
+  (drop (local.get $y))
+ )
+
+ ;; CHECK:      (func $call-various-params-set
+ ;; CHECK-NEXT:  (call $various-params-set
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $various-params-set
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:   (ref.null ${i32_i64})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-set
+  ;; The first argument gets {i32} and {i32}; the second {i32} and {i32_i64;
+  ;; both of those pairs can be optimized to {i32}
+  (call $various-params-set
+   (ref.null ${i32})
+   (ref.null ${i32})
+  )
+  (call $various-params-set
+   (ref.null ${i32})
+   (ref.null ${i32_i64})
+  )
+ )
+ ;; This function is called in ways that *do* allow us to alter the types of
+ ;; its parameters (see last function), however, we reuse the parameters by
+ ;; writing to them, which causes problems in one case.
+ ;; CHECK:      (func $various-params-set (param $x (ref null ${})) (param $y (ref null ${i32}))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (local.set $x
+ ;; CHECK-NEXT:   (ref.null ${})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (local.set $y
+ ;; CHECK-NEXT:   (ref.null ${i32_i64})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-set (param $x (ref null ${})) (param $y (ref null ${}))
+  ;; "Use" the locals to avoid other optimizations kicking in.
+  (drop (local.get $x))
+  (drop (local.get $y))
+  ;; Write to $x in a way that prevents us making the type more specific.
+  (local.set $x (ref.null ${}))
+  ;; Write to $y in a way that still allows us to make the type more specific.
+  (local.set $y (ref.null ${i32_i64}))
+ )
+
+ ;; CHECK:      (func $call-various-params-null
+ ;; CHECK-NEXT:  (call $various-params-null
+ ;; CHECK-NEXT:   (ref.as_non_null
+ ;; CHECK-NEXT:    (ref.null ${i32})
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (ref.null ${i32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $various-params-null
+ ;; CHECK-NEXT:   (ref.as_non_null
+ ;; CHECK-NEXT:    (ref.null ${i32})
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:   (ref.as_non_null
+ ;; CHECK-NEXT:    (ref.null ${i32})
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-null
+  ;; The first argument gets non-null values, allowing us to refine it. The
+  ;; second gets only one.
+  (call $various-params-null
+   (ref.as_non_null (ref.null ${i32}))
+   (ref.null ${i32})
+  )
+  (call $various-params-null
+   (ref.as_non_null (ref.null ${i32}))
+   (ref.as_non_null (ref.null ${i32}))
+  )
+ )
+ ;; This function is called in ways that allow us to make the first parameter
+ ;; non-nullable.
+ ;; CHECK:      (func $various-params-null (param $x (ref ${i32})) (param $y (ref null ${i32}))
+ ;; CHECK-NEXT:  (local $temp i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $y)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (local.set $temp
+ ;; CHECK-NEXT:   (local.get $temp)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-null (param $x (ref null ${})) (param $y (ref null ${}))
+  (local $temp i32)
+  ;; "Use" the locals to avoid other optimizations kicking in.
+  (drop (local.get $x))
+  (drop (local.get $y))
+  ;; Use a local in this function as well, which should be ignored by this pass
+  ;; (when we scan and update all local.gets and sets, we should only do so on
+  ;; parameters, and not vars - and we can crash if we scan/update things we
+  ;; should not).
+  (local.set $temp (local.get $temp))
+ )
+
+ ;; CHECK:      (func $call-various-params-middle
+ ;; CHECK-NEXT:  (call $various-params-middle
+ ;; CHECK-NEXT:   (ref.null ${i32_i64})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (call $various-params-middle
+ ;; CHECK-NEXT:   (ref.null ${i32_f32})
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $call-various-params-middle
+  ;; The argument gets {i32_i64} and {i32_f32}. This allows us to refine from
+  ;; {} to {i32}, a type "in the middle".
+  (call $various-params-middle
+   (ref.null ${i32_i64})
+  )
+  (call $various-params-middle
+   (ref.null ${i32_f32})
+  )
+ )
+ ;; CHECK:      (func $various-params-middle (param $x (ref null ${i32}))
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $x)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $various-params-middle (param $x (ref null ${}))
+  ;; "Use" the local to avoid other optimizations kicking in.
+  (drop (local.get $x))
+ )
 )