Always inline trivial calls that always shrink

src/passes/Inlining.cpp

@@ -67,7 +67,28 @@ enum class InliningMode {
   SplitPatternB
 };
 
-// Useful into on a function, helping us decide if we can inline it
+// Whether a function just calls another function in a way that always shrinks
+// when the calling function is inlined.
+enum class TrivialCall {
+  // Function does not just call another function, or it may not shrink when
+  // inlined.
+  NotTrivial,
+
+  // Function just calls another function, with `local.get`s as arguments, and
+  // with each `local` is used exactly once, and in the order they appear in the
+  // argument list.
+  //
+  // In this case, inlining the function generates smaller code, and it is also
+  // good for runtime.
+  Shrinks,
+
+  // Function just calls another function, but maybe with arguments other than
+  // `local.get`s, or maybe some locals are used more than once. In this case
+  // code size does not always shrink.
+  MayNotShrink,
+};
+
+// Useful info on a function, helping us decide if we can inline it.
 struct FunctionInfo {
   std::atomic<Index> refs;
   Index size;
@@ -77,16 +98,7 @@ struct FunctionInfo {
   // Something is used globally if there is a reference to it in a table or
   // export etc.
   bool usedGlobally;
-  // We consider a function to be a trivial call if the body is just a call with
-  // trivial arguments, like this:
-  //
-  //  (func $forward (param $x) (param $y)
-  //    (call $target (local.get $x) (local.get $y))
-  //  )
-  //
-  // Specifically the body must be a call, and the operands to the call must be
-  // of size 1 (generally, LocalGet or Const).
-  bool isTrivialCall;
+  TrivialCall trivialCall;
   InliningMode inliningMode;
 
   FunctionInfo() { clear(); }
@@ -98,7 +110,7 @@ struct FunctionInfo {
     hasLoops = false;
     hasTryDelegate = false;
     usedGlobally = false;
-    isTrivialCall = false;
+    trivialCall = TrivialCall::NotTrivial;
     inliningMode = InliningMode::Unknown;
   }
 
@@ -110,7 +122,7 @@ struct FunctionInfo {
     hasLoops = other.hasLoops;
     hasTryDelegate = other.hasTryDelegate;
     usedGlobally = other.usedGlobally;
-    isTrivialCall = other.isTrivialCall;
+    trivialCall = other.trivialCall;
     inliningMode = other.inliningMode;
     return *this;
   }
@@ -132,6 +144,11 @@ struct FunctionInfo {
         size <= options.inlining.oneCallerInlineMaxSize) {
       return true;
     }
+    // If the function calls another one in a way that always shrinks when
+    // inlined, inline it in all optimization and shrink modes.
+    if (trivialCall == TrivialCall::Shrinks) {
+      return true;
+    }
     // If it's so big that we have no flexible options that could allow it,
     // do not inline.
     if (size > options.inlining.flexibleInlineMaxSize) {
@@ -143,22 +160,18 @@ struct FunctionInfo {
     if (options.shrinkLevel > 0 || options.optimizeLevel < 3) {
       return false;
     }
-    if (hasCalls) {
-      // This has calls. If it is just a trivial call itself then inline, as we
-      // will save a call that way - basically we skip a trampoline in the
-      // middle - but if it is something more complex, leave it alone, as we may
-      // not help much (and with recursion we may end up with a wasteful
-      // increase in code size).
-      //
-      // Note that inlining trivial calls may increase code size, e.g. if they
-      // use a parameter more than once (forcing us after inlining to save that
-      // value to a local, etc.), but here we are optimizing for speed and not
-      // size, so we risk it.
-      return isTrivialCall;
-    }
-    // This doesn't have calls. Inline if loops do not prevent us (normally, a
-    // loop suggests a lot of work and so inlining is less useful).
-    return !hasLoops || options.inlining.allowFunctionsWithLoops;
+    // The function just calls another function, but it's using locals in
+    // different order than the argument order, and/or using some locals more
+    // than once. In this case we inline if we're not optimizing for code size,
+    // as inlining it to more than one call site may increase code size by
+    // introducing locals.
+    if (trivialCall == TrivialCall::MayNotShrink) {
+      return true;
+    }
+    // Trivial calls are already handled. Inline if
+    // 1. The function doesn't have calls, and
+    // 2. The function doesn't have loops, or we allow inlining with loops.
+    return !hasCalls && (!hasLoops || options.inlining.allowFunctionsWithLoops);
   }
 };
 
@@ -227,10 +240,35 @@ struct FunctionInfoScanner
     info.size = Measurer::measure(curr->body);
 
     if (auto* call = curr->body->dynCast<Call>()) {
+      // If call arguments are function locals read in order, then the code size
+      // always shrinks when the call is inlined. Note that we don't allow
+      // skipping function arguments here, as that can create `drop`
+      // instructions at the call sites, increasing code size.
+      bool shrinks = true;
+      Index nextLocalGetIndex = 0;
+      for (auto* operand : call->operands) {
+        if (auto* localGet = operand->dynCast<LocalGet>()) {
+          if (localGet->index == nextLocalGetIndex) {
+            nextLocalGetIndex += 1;
+          } else {
+            shrinks = false;
+            break;
+          }
+        } else {
+          shrinks = false;
+          break;
+        }
+      }
+
+      if (shrinks) {
+        info.trivialCall = TrivialCall::Shrinks;
+        return;
+      }
+
       if (info.size == call->operands.size() + 1) {
         // This function body is a call with some trivial (size 1) operands like
         // LocalGet or Const, so it is a trivial call.
-        info.isTrivialCall = true;
+        info.trivialCall = TrivialCall::MayNotShrink;
       }
     }
   }

test/lit/passes/inlining-const-args.wat

@@ -0,0 +1,88 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; With `-O3`, we always inline calls to functions that just call other
+;; functions with "trivial" arguments.
+;;
+;; A trivial argument for now is just an instruction with size 1. E.g.
+;; `local.get`, constants.
+
+;; RUN: foreach %s %t wasm-opt -all -O3 -S -o - | filecheck %s --check-prefix=O3
+;; RUN: foreach %s %t wasm-opt -all -O2 -S -o - | filecheck %s --check-prefix=O2
+;; RUN: foreach %s %t wasm-opt -all -Os -S -o - | filecheck %s --check-prefix=Os
+
+(module
+ ;; O3:      (type $0 (func (param i32 i32 i32)))
+ ;; O2:      (type $1 (func))
+
+ ;; O2:      (type $0 (func (param i32 i32 i32)))
+ ;; Os:      (type $1 (func))
+
+ ;; Os:      (type $0 (func (param i32 i32 i32)))
+ (type $0 (func (param i32 i32 i32)))
+
+ ;; O3:      (type $1 (func))
+ (type $1 (func))
+
+ (type $2 (func))
+
+ ;; O3:      (import "env" "foo" (func $imported-foo (type $0) (param i32 i32 i32)))
+ ;; O2:      (import "env" "foo" (func $imported-foo (type $0) (param i32 i32 i32)))
+ ;; Os:      (import "env" "foo" (func $imported-foo (type $0) (param i32 i32 i32)))
+ (import "env" "foo" (func $imported-foo (type $0) (param i32 i32 i32)))
+
+ ;; O3:      (export "main" (func $main))
+ ;; O2:      (export "main" (func $main))
+ ;; Os:      (export "main" (func $main))
+ (export "main" (func $main))
+
+ ;; O2:      (func $call-foo (type $1)
+ ;; O2-NEXT:  (call $imported-foo
+ ;; O2-NEXT:   (i32.const 1)
+ ;; O2-NEXT:   (i32.const 2)
+ ;; O2-NEXT:   (i32.const 3)
+ ;; O2-NEXT:  )
+ ;; O2-NEXT: )
+ ;; Os:      (func $call-foo (type $1)
+ ;; Os-NEXT:  (call $imported-foo
+ ;; Os-NEXT:   (i32.const 1)
+ ;; Os-NEXT:   (i32.const 2)
+ ;; Os-NEXT:   (i32.const 3)
+ ;; Os-NEXT:  )
+ ;; Os-NEXT: )
+ (func $call-foo (type $1)
+  (call $imported-foo
+   (i32.const 1)
+   (i32.const 2)
+   (i32.const 3)))
+
+ ;; O3:      (func $main (type $1)
+ ;; O3-NEXT:  (call $imported-foo
+ ;; O3-NEXT:   (i32.const 1)
+ ;; O3-NEXT:   (i32.const 2)
+ ;; O3-NEXT:   (i32.const 3)
+ ;; O3-NEXT:  )
+ ;; O3-NEXT:  (call $imported-foo
+ ;; O3-NEXT:   (i32.const 1)
+ ;; O3-NEXT:   (i32.const 2)
+ ;; O3-NEXT:   (i32.const 3)
+ ;; O3-NEXT:  )
+ ;; O3-NEXT:  (call $imported-foo
+ ;; O3-NEXT:   (i32.const 1)
+ ;; O3-NEXT:   (i32.const 2)
+ ;; O3-NEXT:   (i32.const 3)
+ ;; O3-NEXT:  )
+ ;; O3-NEXT: )
+ ;; O2:      (func $main (type $1)
+ ;; O2-NEXT:  (call $call-foo)
+ ;; O2-NEXT:  (call $call-foo)
+ ;; O2-NEXT:  (call $call-foo)
+ ;; O2-NEXT: )
+ ;; Os:      (func $main (type $1)
+ ;; Os-NEXT:  (call $call-foo)
+ ;; Os-NEXT:  (call $call-foo)
+ ;; Os-NEXT:  (call $call-foo)
+ ;; Os-NEXT: )
+ (func $main (type $2)
+  (call $call-foo)
+  (call $call-foo)
+  (call $call-foo)))