From 3437bc575de8f34aa7c262519478e00fce8e6df8 Mon Sep 17 00:00:00 2001 From: Alon Zakai Date: Tue, 23 Dec 2025 15:45:52 -0800 Subject: [PATCH 1/5] go --- src/ir/return-utils.cpp | 59 +++++++++++++--------- src/ir/return-utils.h | 7 +++ src/passes/Inlining.cpp | 3 +- test/lit/passes/inlining_splitting.wast | 65 +++++++++++++++++++++++++ 4 files changed, 111 insertions(+), 23 deletions(-) diff --git a/src/ir/return-utils.cpp b/src/ir/return-utils.cpp index 20b3a194b13..5da36c8deea 100644 --- a/src/ir/return-utils.cpp +++ b/src/ir/return-utils.cpp @@ -56,6 +56,42 @@ void removeReturns(Function* func, Module& wasm) { ReturnValueRemover().walkFunctionInModule(func, &wasm); } +Info getInfo(Expression* curr) { + struct Finder : PostWalker { + Info info; + + Finder(Info& info) : info(info) {} + + void visitReturn(Return* curr) { + info.hasReturn = true; + } + + void visitCall(Call* curr) { + if (curr->isReturn) { + info.hasReturnCall = true; + } + } + void visitCallIndirect(CallIndirect* curr) { + if (curr->isReturn) { + info.hasReturnCall = true; + } + } + void visitCallRef(CallRef* curr) { + if (curr->isReturn) { + info.hasReturnCall = true; + } + } + } finder; + finder.walk(curr); + + // Any return call is also a return. + if (finder.info.hasReturnCall) { + finder.info.hasReturn = true; + } + + return finder.info; +} + std::unordered_map findReturnCallers(Module& wasm) { ModuleUtils::ParallelFunctionAnalysis analysis( wasm, [&](Function* func, bool& hasReturnCall) { @@ -63,28 +99,7 @@ std::unordered_map findReturnCallers(Module& wasm) { return; } - struct Finder : PostWalker { - bool hasReturnCall = false; - - void visitCall(Call* curr) { - if (curr->isReturn) { - hasReturnCall = true; - } - } - void visitCallIndirect(CallIndirect* curr) { - if (curr->isReturn) { - hasReturnCall = true; - } - } - void visitCallRef(CallRef* curr) { - if (curr->isReturn) { - hasReturnCall = true; - } - } - } finder; - - finder.walk(func->body); - hasReturnCall = finder.hasReturnCall; + hasReturnCall = getInfo(func->body).hasReturnCall; }); // Convert to an unordered map for fast lookups. TODO: Avoid a copy here. diff --git a/src/ir/return-utils.h b/src/ir/return-utils.h index a5214ba017c..ce1169c321f 100644 --- a/src/ir/return-utils.h +++ b/src/ir/return-utils.h @@ -34,6 +34,13 @@ void removeReturns(Function* func, Module& wasm); using ReturnCallersMap = std::unordered_map; ReturnCallersMap findReturnCallers(Module& wasm); +// Information about the returns in a given expression. +struct Info { + bool hasReturn = false; + bool hasReturnCall = false; +}; +Info getInfo(Expression* curr); + } // namespace wasm::ReturnUtils #endif // wasm_ir_return_h diff --git a/src/passes/Inlining.cpp b/src/passes/Inlining.cpp index 98be0a7636d..41af04e96ba 100644 --- a/src/passes/Inlining.cpp +++ b/src/passes/Inlining.cpp @@ -42,6 +42,7 @@ #include "ir/module-utils.h" #include "ir/names.h" #include "ir/properties.h" +#include "ir/return-utils.h" #include "ir/type-updating.h" #include "ir/utils.h" #include "parsing.h" @@ -1007,7 +1008,7 @@ struct FunctionSplitter { } if (iff->ifTrue->type == Type::none) { // This must have no returns. - if (!FindAll(iff->ifTrue).list.empty()) { + if (ReturnUtils::getInfo(iff->ifTrue).hasReturn) { return InliningMode::Uninlineable; } } else { diff --git a/test/lit/passes/inlining_splitting.wast b/test/lit/passes/inlining_splitting.wast index e583e2e4021..5c9630145a6 100644 --- a/test/lit/passes/inlining_splitting.wast +++ b/test/lit/passes/inlining_splitting.wast @@ -1276,6 +1276,71 @@ (call $unreachable-if-body-no-result (ref.null any)) ) + + + + (func $nop (param $x anyref) (result anyref) + ;; A loop and recursion, to avoid this getting inlined (we want to keep the + ;; return_call to here as a return_call). + (loop $loop + (if + (ref.is_null + (local.get $x) + ) + (then + (br $loop) + ) + ) + (return_call $nop + (local.get $x) + ) + ) + ) + + (func $reachable-if-body-return_call (param $x anyref) (result anyref) + (if + (ref.is_null + (local.get $x) + ) + (then + (if + (i32.const 1) + ;; The return_call here prevents the optimization. + (then + (return_call $nop + (local.get $x) + ) + ) + (else + (call $import) + ) + ) + ) + ) + (local.get $x) + ) + + ;; CHECK: (func $call-reachable-if-body-return_call (type $0) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (call $reachable-if-body-return_call + ;; CHECK-NEXT: (ref.null none) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (drop + ;; CHECK-NEXT: (call $reachable-if-body-return_call + ;; CHECK-NEXT: (ref.null none) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $call-reachable-if-body-return_call + (drop (call $reachable-if-body-return_call (ref.null any))) + (drop (call $reachable-if-body-return_call (ref.null any))) + ) + + + + + (func $multi-if (param $x anyref) (result anyref) (if (ref.is_null From 9aa80e87848d3b5d4b23b1d869fd6280278163cb Mon Sep 17 00:00:00 2001 From: Alon Zakai Date: Tue, 23 Dec 2025 15:46:11 -0800 Subject: [PATCH 2/5] go --- src/ir/return-utils.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/ir/return-utils.cpp b/src/ir/return-utils.cpp index 5da36c8deea..3069d3bead3 100644 --- a/src/ir/return-utils.cpp +++ b/src/ir/return-utils.cpp @@ -60,8 +60,6 @@ Info getInfo(Expression* curr) { struct Finder : PostWalker { Info info; - Finder(Info& info) : info(info) {} - void visitReturn(Return* curr) { info.hasReturn = true; } From 9e9d1f8e7f58e7f4ba9fc962a057023a42736ef4 Mon Sep 17 00:00:00 2001 From: Alon Zakai Date: Tue, 23 Dec 2025 15:46:54 -0800 Subject: [PATCH 3/5] finish --- test/lit/passes/inlining_splitting.wast | 48 +++++++++++++++++++++---- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/test/lit/passes/inlining_splitting.wast b/test/lit/passes/inlining_splitting.wast index 5c9630145a6..e81cf4702ae 100644 --- a/test/lit/passes/inlining_splitting.wast +++ b/test/lit/passes/inlining_splitting.wast @@ -828,7 +828,7 @@ ;; CHECK-NEXT: (local.get $0) ;; CHECK-NEXT: ) ;; CHECK-NEXT: (then - ;; CHECK-NEXT: (call $byn-split-outlined-A$colliding-name_67 + ;; CHECK-NEXT: (call $byn-split-outlined-A$colliding-name_70 ;; CHECK-NEXT: (local.get $0) ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) @@ -843,7 +843,7 @@ ;; CHECK-NEXT: (local.get $1) ;; CHECK-NEXT: ) ;; CHECK-NEXT: (then - ;; CHECK-NEXT: (call $byn-split-outlined-A$colliding-name_67 + ;; CHECK-NEXT: (call $byn-split-outlined-A$colliding-name_70 ;; CHECK-NEXT: (local.get $1) ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) @@ -1279,6 +1279,21 @@ + ;; CHECK: (func $nop (type $2) (param $x anyref) (result anyref) + ;; CHECK-NEXT: (loop $loop + ;; CHECK-NEXT: (if + ;; CHECK-NEXT: (ref.is_null + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (br $loop) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (return_call $nop + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) (func $nop (param $x anyref) (result anyref) ;; A loop and recursion, to avoid this getting inlined (we want to keep the ;; return_call to here as a return_call). @@ -1297,6 +1312,27 @@ ) ) + ;; CHECK: (func $reachable-if-body-return_call (type $2) (param $x anyref) (result anyref) + ;; CHECK-NEXT: (if + ;; CHECK-NEXT: (ref.is_null + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (if + ;; CHECK-NEXT: (i32.const 1) + ;; CHECK-NEXT: (then + ;; CHECK-NEXT: (return_call $nop + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (else + ;; CHECK-NEXT: (call $import) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) (func $reachable-if-body-return_call (param $x anyref) (result anyref) (if (ref.is_null @@ -1396,7 +1432,7 @@ ;; CHECK-NEXT: (local.get $0) ;; CHECK-NEXT: ) ;; CHECK-NEXT: (then - ;; CHECK-NEXT: (call $byn-split-outlined-B$multi-if_76 + ;; CHECK-NEXT: (call $byn-split-outlined-B$multi-if_79 ;; CHECK-NEXT: (local.get $0) ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) @@ -1429,7 +1465,7 @@ ;; CHECK-NEXT: (local.get $1) ;; CHECK-NEXT: ) ;; CHECK-NEXT: (then - ;; CHECK-NEXT: (call $byn-split-outlined-B$multi-if_76 + ;; CHECK-NEXT: (call $byn-split-outlined-B$multi-if_79 ;; CHECK-NEXT: (local.get $1) ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) @@ -1601,7 +1637,7 @@ ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) -;; CHECK: (func $byn-split-outlined-A$colliding-name_67 (type $1) (param $x i32) +;; CHECK: (func $byn-split-outlined-A$colliding-name_70 (type $1) (param $x i32) ;; CHECK-NEXT: (loop $l ;; CHECK-NEXT: (call $import) ;; CHECK-NEXT: (br $l) @@ -1624,7 +1660,7 @@ ;; CHECK-NEXT: (unreachable) ;; CHECK-NEXT: ) -;; CHECK: (func $byn-split-outlined-B$multi-if_76 (type $3) (param $x anyref) +;; CHECK: (func $byn-split-outlined-B$multi-if_79 (type $3) (param $x anyref) ;; CHECK-NEXT: (loop $x ;; CHECK-NEXT: (call $import) ;; CHECK-NEXT: (br_if $x From 33b327c9cac205e0dc7884368d7500a49ad3b0d7 Mon Sep 17 00:00:00 2001 From: Alon Zakai Date: Tue, 23 Dec 2025 15:46:59 -0800 Subject: [PATCH 4/5] form --- src/ir/return-utils.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ir/return-utils.cpp b/src/ir/return-utils.cpp index 3069d3bead3..ed21cde4863 100644 --- a/src/ir/return-utils.cpp +++ b/src/ir/return-utils.cpp @@ -60,9 +60,7 @@ Info getInfo(Expression* curr) { struct Finder : PostWalker { Info info; - void visitReturn(Return* curr) { - info.hasReturn = true; - } + void visitReturn(Return* curr) { info.hasReturn = true; } void visitCall(Call* curr) { if (curr->isReturn) { From 6b7078eba77113bbe01d01d5927057d7c3a3a56f Mon Sep 17 00:00:00 2001 From: Alon Zakai Date: Tue, 23 Dec 2025 15:49:16 -0800 Subject: [PATCH 5/5] nice --- test/lit/passes/inlining_splitting.wast | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/test/lit/passes/inlining_splitting.wast b/test/lit/passes/inlining_splitting.wast index e81cf4702ae..fcb1b45c5e6 100644 --- a/test/lit/passes/inlining_splitting.wast +++ b/test/lit/passes/inlining_splitting.wast @@ -1276,9 +1276,6 @@ (call $unreachable-if-body-no-result (ref.null any)) ) - - - ;; CHECK: (func $nop (type $2) (param $x anyref) (result anyref) ;; CHECK-NEXT: (loop $loop ;; CHECK-NEXT: (if @@ -1341,7 +1338,9 @@ (then (if (i32.const 1) - ;; The return_call here prevents the optimization. + ;; The return_call here prevents the optimization (nothing is inlined + ;; in our caller, below), just like the return in the previous + ;; testcase. (then (return_call $nop (local.get $x) @@ -1369,14 +1368,11 @@ ;; CHECK-NEXT: ) ;; CHECK-NEXT: ) (func $call-reachable-if-body-return_call + ;; We do not inline here. (drop (call $reachable-if-body-return_call (ref.null any))) (drop (call $reachable-if-body-return_call (ref.null any))) ) - - - - (func $multi-if (param $x anyref) (result anyref) (if (ref.is_null