feat: implement tail call optimization

src/passes/CMakeLists.txt

@@ -116,6 +116,7 @@ set(passes_SOURCES
   ReorderGlobals.cpp
   ReorderLocals.cpp
   ReReloop.cpp
+  TailCall.cpp
   TrapMode.cpp
   TypeGeneralizing.cpp
   TypeRefining.cpp

src/passes/TailCall.cpp

@@ -0,0 +1,214 @@
+
+#include "cfg/cfg-traversal.h"
+#include "ir/properties.h"
+#include "ir/utils.h"
+#include "pass.h"
+#include "wasm-traversal.h"
+#include "wasm.h"
+#include <algorithm>
+#include <cassert>
+#include <iostream>
+#include <optional>
+#include <stack>
+#include <vector>
+
+namespace wasm {
+
+namespace {
+
+struct Info {
+  bool isStartWithReturn = false;
+  bool isInsideTryBlock = false;
+  Expression* lastExpr = nullptr;
+};
+
+struct NonReturnFinder
+  : public CFGWalker<NonReturnFinder,
+                     UnifiedExpressionVisitor<NonReturnFinder>,
+                     Info> {
+  using S =
+    CFGWalker<NonReturnFinder, UnifiedExpressionVisitor<NonReturnFinder>, Info>;
+
+  std::vector<Call*> tailCalls;
+  std::vector<CallIndirect*> tailCallIndirects;
+
+  void visitExpression(Expression* curr) {
+    if (currBasicBlock == nullptr) {
+      return;
+    }
+    if (curr->is<Block>() || curr->is<If>() || curr->is<Loop>()) {
+      // skip all control flow instructions
+      return;
+    }
+
+    Expression* const lastExpr = currBasicBlock->contents.lastExpr;
+    currBasicBlock->contents.lastExpr = curr;
+
+    if (!tryStack.empty()) {
+      // skip all try stack
+      currBasicBlock->contents.isInsideTryBlock = true;
+    }
+    if (curr->is<Return>()) {
+      if (lastExpr == nullptr) {
+        currBasicBlock->contents.isStartWithReturn = true;
+      } else {
+        pushPotentialTailCall(lastExpr);
+      }
+    }
+  }
+
+  void pushPotentialTailCall(Expression* curr) {
+    if (curr) {
+      if (curr->is<Call>()) {
+        tailCalls.push_back(curr->cast<Call>());
+      } else if (curr->is<CallIndirect>()) {
+        tailCallIndirects.push_back(curr->cast<CallIndirect>());
+      }
+    }
+  }
+
+  void doWalkFunction(Function* func) {
+    S::doWalkFunction(func);
+    if (hasSyntheticExit && exit != nullptr) {
+      exit->contents.isStartWithReturn = true;
+    }
+    if (exit != nullptr) {
+      assert(tryStack.empty());
+      pushPotentialTailCall(exit->contents.lastExpr);
+    }
+    // propagate start with return flag
+    bool hasUpdated = true;
+    while (hasUpdated) {
+      hasUpdated = false;
+      for (std::unique_ptr<BasicBlock> const& bb : basicBlocks) {
+        if (bb->contents.isStartWithReturn) {
+          continue;
+        }
+        if (bb->contents.lastExpr == nullptr) {
+          const bool followBasicBlockStartWithReturn =
+            std::all_of(bb->out.begin(), bb->out.end(), [](BasicBlock* b) {
+              return b->contents.isStartWithReturn;
+            });
+          if (followBasicBlockStartWithReturn) {
+            bb->contents.isStartWithReturn = true;
+            hasUpdated = true;
+          }
+        }
+      }
+    }
+    for (std::unique_ptr<BasicBlock> const& bb : basicBlocks) {
+      Expression* const lastExpr = bb->contents.lastExpr;
+      if (lastExpr == nullptr) {
+        continue;
+      }
+      const bool followBasicBlockStartWithReturn =
+        std::all_of(bb->out.begin(), bb->out.end(), [](BasicBlock* b) {
+          return b->contents.isStartWithReturn;
+        });
+      if (!followBasicBlockStartWithReturn) {
+        continue;
+      }
+      pushPotentialTailCall(lastExpr);
+    }
+  }
+};
+
+struct ReturnFinder : TryDepthWalker<ReturnFinder> {
+  explicit ReturnFinder(const PassOptions& passOptions)
+    : TryDepthWalker<ReturnFinder>(), passOptions(passOptions) {}
+  const PassOptions& passOptions;
+  std::vector<Call*> tailCalls;
+  std::vector<CallIndirect*> tailCallIndirects;
+  void visitFunction(Function* curr) { checkTailCall(curr->body); }
+  void visitReturn(Return* curr) {
+    if (tryDepth > 0) {
+      // (return (call ...)) is not equal to (return_call ...) in try block
+      return;
+    }
+    checkTailCall(curr->value);
+  }
+
+private:
+  void checkTailCall(Expression* const curr) {
+    std::stack<Expression*> workList{};
+    workList.push(curr);
+    while (!workList.empty()) {
+      Expression* const target = workList.top();
+      workList.pop();
+      if (auto* call = target->dynCast<Call>()) {
+        if (!call->isReturn && call->type == getFunction()->getResults()) {
+          tailCalls.push_back(call);
+        }
+      } else if (auto* call = target->dynCast<CallIndirect>()) {
+        if (!call->isReturn && call->type == getFunction()->getResults()) {
+          tailCallIndirects.push_back(call);
+        }
+      } else if (auto* ifElse = target->dynCast<If>()) {
+        workList.push(ifElse->ifTrue);
+        workList.push(ifElse->ifFalse);
+      } else if (auto* tryy = target->dynCast<Try>()) {
+        for (Expression* catchBody : tryy->catchBodies) {
+          workList.push(catchBody);
+        }
+      } else if (auto* block = target->dynCast<Block>()) {
+        if (!block->list.empty()) {
+          workList.push(block->list.back());
+        }
+      } else {
+        Expression* const next = Properties::getImmediateFallthrough(
+          target, passOptions, *getModule());
+        if (next != target) {
+          workList.push(next);
+        }
+      }
+    }
+  }
+};
+
+} // namespace
+
+struct TailCallOptimizer : public Pass {
+  bool isFunctionParallel() override { return true; }
+  std::unique_ptr<Pass> create() override {
+    return std::make_unique<TailCallOptimizer>();
+  }
+
+  static void modify(std::vector<Call*> const& tailCalls,
+                     std::vector<CallIndirect*> const& tailCallIndirects) {
+    for (Call* call : tailCalls) {
+      if (!call->isReturn) {
+        call->isReturn = true;
+      }
+    }
+    for (CallIndirect* call : tailCallIndirects) {
+      if (!call->isReturn) {
+        call->isReturn = true;
+      }
+    }
+  }
+  void runOnFunction(Module* module, Function* function) override {
+    if (!module->features.hasTailCall()) {
+      return;
+    }
+    if (getPassOptions().shrinkLevel > 0 &&
+        getPassOptions().optimizeLevel == 0) {
+      // When we more force on the binary size, add return_call will increase
+      // the code size.
+      return;
+    }
+    if (function->getResults().size() == 0) {
+      NonReturnFinder finder{};
+      finder.walkFunctionInModule(function, module);
+      modify(finder.tailCalls, finder.tailCallIndirects);
+    } else {
+      ReturnFinder finder{getPassOptions()};
+      finder.walkFunctionInModule(function, module);
+      modify(finder.tailCalls, finder.tailCallIndirects);
+    }
+    ReFinalize{}.walkFunctionInModule(function, module);
+  }
+};
+
+Pass* createTailCallPass() { return new TailCallOptimizer(); }
+
+} // namespace wasm

src/passes/pass.cpp

@@ -552,6 +552,9 @@ void PassRegistry::registerPasses() {
   registerPass("strip-target-features",
                "strip the wasm target features section",
                createStripTargetFeaturesPass);
+  registerPass("tail-call-optimization",
+               "transform call to return call",
+               createTailCallPass);
   registerPass("translate-to-new-eh",
                "deprecated; same as translate-to-exnref",
                createTranslateToExnrefPass);

src/passes/passes.h

@@ -180,6 +180,7 @@ Pass* createStripEHPass();
 Pass* createStubUnsupportedJSOpsPass();
 Pass* createSSAifyPass();
 Pass* createSSAifyNoMergePass();
+Pass* createTailCallPass();
 Pass* createTable64LoweringPass();
 Pass* createTranslateToExnrefPass();
 Pass* createTrapModeClamp();

test/lit/help/wasm-metadce.test

@@ -530,6 +530,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/help/wasm-opt.test

@@ -554,6 +554,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/help/wasm2js.test

@@ -494,6 +494,8 @@
 ;; CHECK-NEXT:
 ;; CHECK-NEXT:   --table64-lowering                            alias for memory64-lowering
 ;; CHECK-NEXT:
+;; CHECK-NEXT:   --tail-call-optimization                      transform call to return call 
+;; CHECK-NEXT:
 ;; CHECK-NEXT:   --trace-calls                                 instrument the build with code
 ;; CHECK-NEXT:                                                 to intercept specific function
 ;; CHECK-NEXT:                                                 calls

test/lit/tail-call-optimization-eh.wast

@@ -0,0 +1,94 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --tail-call-optimization --enable-tail-call --enable-exception-handling --optimize-level 2 --shrink-level 0 -S -o - | filecheck %s
+
+;; Tests for tail call optimization with exception handling
+
+(module $exception
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (type $1 (func))
+
+  ;; CHECK:      (tag $empty (type $1))
+  (tag $empty)
+  ;; CHECK:      (func $f (result i32)
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $f (result i32)
+    i32.const 0
+  )
+  ;; CHECK:      (func $in_try (result i32)
+  ;; CHECK-NEXT:  (try (result i32)
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $in_try (result i32)
+    try (result i32)
+      call $f
+    catch $empty
+      call $f
+    end
+  )
+  ;; CHECK:      (func $out_try (result i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (return_call $f)
+  ;; CHECK-NEXT: )
+  (func $out_try (result i32)
+    try
+    catch $empty
+    end
+    call $f
+  )
+  ;; CHECK:      (func $in_catch (result i32)
+  ;; CHECK-NEXT:  (try
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return
+  ;; CHECK-NEXT:     (return_call $f)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $in_catch (result i32)
+    try
+    catch $empty
+      call $f
+      return
+    end
+    i32.const 0
+  )
+  ;; CHECK:      (func $implicit_in_catch (result i32)
+  ;; CHECK-NEXT:  (try (result i32)
+  ;; CHECK-NEXT:   (do
+  ;; CHECK-NEXT:    (i32.const 0)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch $empty
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (catch_all
+  ;; CHECK-NEXT:    (return_call $f)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $implicit_in_catch (result i32)
+    try (result i32)
+      i32.const 0
+    catch $empty
+      call $f
+    catch_all
+      call $f
+    end
+  )
+)

test/lit/tail-call-optimization-shrink.wast

@@ -0,0 +1,22 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: foreach %s %t wasm-opt --tail-call-optimization --enable-tail-call --optimize-level 0 --shrink-level 2 -S -o - | filecheck %s
+
+;; Tests for tail call optimization
+
+(module
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (func $f (result i32)
+  ;; CHECK-NEXT:  (i32.const 0)
+  ;; CHECK-NEXT: )
+  (func $f (result i32)
+    i32.const 0
+  )
+  ;; CHECK:      (func $implicit_return (result i32)
+  ;; CHECK-NEXT:  (call $f)
+  ;; CHECK-NEXT: )
+  (func $implicit_return (result i32)
+    call $f
+  )
+)