diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt index 16891caade3..96907e9565d 100644 --- a/src/passes/CMakeLists.txt +++ b/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 diff --git a/src/passes/TailCall.cpp b/src/passes/TailCall.cpp new file mode 100644 index 00000000000..88ae1d37121 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +namespace wasm { + +namespace { + +struct Info { + bool isStartWithReturn = false; + bool isInsideTryBlock = false; + Expression* lastExpr = nullptr; +}; + +struct NonReturnFinder + : public CFGWalker, + Info> { + using S = + CFGWalker, Info>; + + std::vector tailCalls; + std::vector tailCallIndirects; + + void visitExpression(Expression* curr) { + if (currBasicBlock == nullptr) { + return; + } + if (curr->is() || curr->is() || curr->is()) { + // 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()) { + if (lastExpr == nullptr) { + currBasicBlock->contents.isStartWithReturn = true; + } else { + pushPotentialTailCall(lastExpr); + } + } + } + + void pushPotentialTailCall(Expression* curr) { + if (curr) { + if (curr->is()) { + tailCalls.push_back(curr->cast()); + } else if (curr->is()) { + tailCallIndirects.push_back(curr->cast()); + } + } + } + + 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 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 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 { + explicit ReturnFinder(const PassOptions& passOptions) + : TryDepthWalker(), passOptions(passOptions) {} + const PassOptions& passOptions; + std::vector tailCalls; + std::vector 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 workList{}; + workList.push(curr); + while (!workList.empty()) { + Expression* const target = workList.top(); + workList.pop(); + if (auto* call = target->dynCast()) { + if (!call->isReturn && call->type == getFunction()->getResults()) { + tailCalls.push_back(call); + } + } else if (auto* call = target->dynCast()) { + if (!call->isReturn && call->type == getFunction()->getResults()) { + tailCallIndirects.push_back(call); + } + } else if (auto* ifElse = target->dynCast()) { + workList.push(ifElse->ifTrue); + workList.push(ifElse->ifFalse); + } else if (auto* tryy = target->dynCast()) { + for (Expression* catchBody : tryy->catchBodies) { + workList.push(catchBody); + } + } else if (auto* block = target->dynCast()) { + 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 create() override { + return std::make_unique(); + } + + static void modify(std::vector const& tailCalls, + std::vector 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 diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp index 8db59b8ac0f..a2689c4091c 100644 --- a/src/passes/pass.cpp +++ b/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); diff --git a/src/passes/passes.h b/src/passes/passes.h index 92dcd3e4eb4..b7d8033d700 100644 --- a/src/passes/passes.h +++ b/src/passes/passes.h @@ -180,6 +180,7 @@ Pass* createStripEHPass(); Pass* createStubUnsupportedJSOpsPass(); Pass* createSSAifyPass(); Pass* createSSAifyNoMergePass(); +Pass* createTailCallPass(); Pass* createTable64LoweringPass(); Pass* createTranslateToExnrefPass(); Pass* createTrapModeClamp(); diff --git a/test/lit/help/wasm-metadce.test b/test/lit/help/wasm-metadce.test index 4c727bcbd1f..7bb5951920f 100644 --- a/test/lit/help/wasm-metadce.test +++ b/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 diff --git a/test/lit/help/wasm-opt.test b/test/lit/help/wasm-opt.test index a0d8d199f94..d546114945c 100644 --- a/test/lit/help/wasm-opt.test +++ b/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 diff --git a/test/lit/help/wasm2js.test b/test/lit/help/wasm2js.test index 881e18950e7..f62457ead1c 100644 --- a/test/lit/help/wasm2js.test +++ b/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 diff --git a/test/lit/tail-call-optimization-eh.wast b/test/lit/tail-call-optimization-eh.wast new file mode 100644 index 00000000000..13e1f0b96d3 --- /dev/null +++ b/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 + ) +) diff --git a/test/lit/tail-call-optimization-shrink.wast b/test/lit/tail-call-optimization-shrink.wast new file mode 100644 index 00000000000..5fca8e8b76b --- /dev/null +++ b/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 + ) +) diff --git a/test/lit/tail-call-optimization.wast b/test/lit/tail-call-optimization.wast new file mode 100644 index 00000000000..282261c97e0 --- /dev/null +++ b/test/lit/tail-call-optimization.wast @@ -0,0 +1,202 @@ +;; 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 2 --shrink-level 0 -S -o - | filecheck %s + +;; Tests for tail call optimization + +(module + ;; CHECK: (type $0 (func (result i32))) + + ;; CHECK: (type $1 (func (param i32) (result i32))) + + ;; CHECK: (type $2 (func (param i32))) + + ;; CHECK: (type $3 (func)) + + ;; 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: (return_call $f) + ;; CHECK-NEXT: ) + (func $implicit_return (result i32) + call $f + ) + ;; CHECK: (func $explicit_return (result i32) + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + (func $explicit_return (result i32) + call $f + ) + ;; CHECK: (func $return_through_tee (param $0 i32) (result i32) + ;; CHECK-NEXT: (local.tee $0 + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_tee (param $0 i32) (result i32) + call $f + local.tee $0 + ) + ;; CHECK: (func $return_through_block (result i32) + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + (func $return_through_block (result i32) + block (result i32) + call $f + end + ) + ;; CHECK: (func $return_through_loop (result i32) + ;; CHECK-NEXT: (loop + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_loop (result i32) + loop (result i32) + call $f + end + ) + ;; CHECK: (func $return_through_if_then (param $0 i32) (result i32) + ;; CHECK-NEXT: (if (result i32) + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (else + ;; CHECK-NEXT: (i32.const 0) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_if_then (param $0 i32) (result i32) + local.get $0 + if (result i32) + call $f + else + i32.const 0 + end + ) + ;; CHECK: (func $return_through_if_else (param $0 i32) (result i32) + ;; CHECK-NEXT: (if (result i32) + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (i32.const 0) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (else + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_if_else (param $0 i32) (result i32) + local.get $0 + if (result i32) + i32.const 0 + else + call $f + end + ) + ;; CHECK: (func $return_through_if_both (param $0 i32) (result i32) + ;; CHECK-NEXT: (if + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (else + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_if_both (param $0 i32) (result i32) + local.get $0 + if (result i32) + call $f + else + call $f + end + ) + ;; CHECK: (func $return_through_br_if (param $0 i32) (result i32) + ;; CHECK-NEXT: (block $block + ;; CHECK-NEXT: (block + ;; CHECK-NEXT: (return_call $f) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (i32.const 1) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_br_if (param $0 i32) (result i32) + block (result i32) + call $f + i32.const 1 + br_if 0 + end + ) + + ;; CHECK: (func $g (param $0 i32) + ;; CHECK-NEXT: ) + (func $g (param i32)) + ;; CHECK: (func $return_without_value + ;; CHECK-NEXT: (return_call $g + ;; CHECK-NEXT: (i32.const 1) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_without_value + i32.const 1 + call $g + ) + ;; CHECK: (func $return_without_value_through_block (param $0 i32) + ;; CHECK-NEXT: (return_call $g + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_without_value_through_block (param $0 i32) + block + local.get 0 + call $g + end + ) + ;; CHECK: (func $return_without_value_through_if (param $0 i32) + ;; CHECK-NEXT: (if + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (return_call $g + ;; CHECK-NEXT: (local.get $0) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_without_value_through_if (param $0 i32) + local.get 0 + if + local.get 0 + call $g + end + ) +) + +(module $NYI + ;; CHECK: (type $0 (func (result i32))) + + ;; CHECK: (type $1 (func (param i32) (result i32))) + + ;; CHECK: (func $f (result i32) + ;; CHECK-NEXT: (i32.const 0) + ;; CHECK-NEXT: ) + (func $f (result i32) + i32.const 0 + ) + ;; CHECK: (func $return_through_br (param $0 i32) (result i32) + ;; CHECK-NEXT: (block $block (result i32) + ;; CHECK-NEXT: (br $block + ;; CHECK-NEXT: (call $f) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $return_through_br (param $0 i32) (result i32) + block (result i32) + call $f + br 0 + end + ) +)