diff options
| -rw-r--r-- | src/ir/module-splitting.cpp | 66 | ||||
| -rw-r--r-- | test/lit/wasm-split/segments.wast | 94 |
2 files changed, 156 insertions, 4 deletions
diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp index 8ef0e7b8c..27b6a0902 100644 --- a/src/ir/module-splitting.cpp +++ b/src/ir/module-splitting.cpp @@ -282,7 +282,7 @@ struct ModuleSplitter { // Initialization helpers static std::unique_ptr<Module> initSecondary(const Module& primary); static std::pair<std::set<Name>, std::set<Name>> - classifyFunctions(const Module& primary, const Config& config); + classifyFunctions(Module& primary, const Config& config); static std::map<Name, Name> initExportedPrimaryFuncs(const Module& primary); // Other helpers @@ -343,7 +343,64 @@ std::unique_ptr<Module> ModuleSplitter::initSecondary(const Module& primary) { } std::pair<std::set<Name>, std::set<Name>> -ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) { +ModuleSplitter::classifyFunctions(Module& primary, const Config& config) { + // Find functions that refer to data or element segments. These functions must + // remain in the primary module because segments cannot be exported to be + // accessed from the secondary module. + // + // TODO: Investigate other options, such as moving the segments to the + // secondary module or replacing the segment-using instructions in the + // secondary module with calls to imports. + ModuleUtils::ParallelFunctionAnalysis<std::vector<Name>> + segmentReferrerCollector( + primary, [&](Function* func, std::vector<Name>& segmentReferrers) { + if (func->imported()) { + return; + } + + struct SegmentReferrerCollector + : PostWalker<SegmentReferrerCollector, + UnifiedExpressionVisitor<SegmentReferrerCollector>> { + bool hasSegmentReference = false; + + void visitExpression(Expression* curr) { + +#define DELEGATE_ID curr->_id + +#define DELEGATE_START(id) [[maybe_unused]] auto* cast = curr->cast<id>(); +#define DELEGATE_GET_FIELD(id, field) cast->field +#define DELEGATE_FIELD_TYPE(id, field) +#define DELEGATE_FIELD_HEAPTYPE(id, field) +#define DELEGATE_FIELD_CHILD(id, field) +#define DELEGATE_FIELD_OPTIONAL_CHILD(id, field) +#define DELEGATE_FIELD_INT(id, field) +#define DELEGATE_FIELD_LITERAL(id, field) +#define DELEGATE_FIELD_NAME(id, field) +#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field) +#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field) +#define DELEGATE_FIELD_ADDRESS(id, field) + +#define DELEGATE_FIELD_NAME_KIND(id, field, kind) \ + if (kind == ModuleItemKind::DataSegment || \ + kind == ModuleItemKind::ElementSegment) { \ + hasSegmentReference = true; \ + } + +#include "wasm-delegations-fields.def" + } + }; + SegmentReferrerCollector collector; + collector.walkFunction(func); + if (collector.hasSegmentReference) { + segmentReferrers.push_back(func->name); + } + }); + + std::unordered_set<Name> segmentReferrers; + for (auto& [_, referrers] : segmentReferrerCollector.map) { + segmentReferrers.insert(referrers.begin(), referrers.end()); + } + std::set<Name> primaryFuncs, secondaryFuncs; for (auto& func : primary.functions) { // In JSPI mode exported functions cannot be moved to the secondary @@ -351,14 +408,15 @@ ModuleSplitter::classifyFunctions(const Module& primary, const Config& config) { // wrapper. Exported JSPI functions can still benefit from splitting though // since only the JSPI wrapper stub will remain in the primary module. if (func->imported() || config.primaryFuncs.count(func->name) || - (config.jspi && ExportUtils::isExported(primary, *func))) { + (config.jspi && ExportUtils::isExported(primary, *func)) || + segmentReferrers.count(func->name)) { primaryFuncs.insert(func->name); } else { assert(func->name != primary.start && "The start function must be kept"); secondaryFuncs.insert(func->name); } } - return std::make_pair(primaryFuncs, secondaryFuncs); + return std::make_pair(std::move(primaryFuncs), std::move(secondaryFuncs)); } std::map<Name, Name> diff --git a/test/lit/wasm-split/segments.wast b/test/lit/wasm-split/segments.wast new file mode 100644 index 000000000..7a14afc42 --- /dev/null +++ b/test/lit/wasm-split/segments.wast @@ -0,0 +1,94 @@ +;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited. + +;; RUN: wasm-split %s -all --keep-funcs=foo -g -o1 %t.1.wasm -o2 %t.2.wasm +;; RUN: wasm-dis %t.1.wasm | filecheck %s --check-prefix PRIMARY +;; RUN: wasm-dis %t.2.wasm | filecheck %s --check-prefix SECONDARY + +(module + + ;; PRIMARY: (type $0 (func)) + + ;; PRIMARY: (type $data-array (array i8)) + (type $data-array (array i8)) + + ;; PRIMARY: (type $elem-array (array externref)) + (type $elem-array (array externref)) + + ;; PRIMARY: (memory $mem 0) + (memory $mem 0) + + ;; PRIMARY: (data $data "hello world") + (data $data "hello world") + + ;; PRIMARY: (elem $elem externref) + (elem $elem externref) + + ;; PRIMARY: (export "memory" (memory $mem)) + + ;; PRIMARY: (func $data.drop + ;; PRIMARY-NEXT: (data.drop $data) + ;; PRIMARY-NEXT: ) + (func $data.drop + (data.drop $data) + ) + + ;; PRIMARY: (func $memory.init + ;; PRIMARY-NEXT: (memory.init $data + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: ) + ;; PRIMARY-NEXT: ) + (func $memory.init + (memory.init $mem $data + (i32.const 0) + (i32.const 0) + (i32.const 0) + ) + ) + + ;; PRIMARY: (func $array.new_data + ;; PRIMARY-NEXT: (drop + ;; PRIMARY-NEXT: (array.new_data $data-array $data + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: ) + ;; PRIMARY-NEXT: ) + ;; PRIMARY-NEXT: ) + (func $array.new_data + (drop + (array.new_data $data-array $data + (i32.const 0) + (i32.const 0) + ) + ) + ) + + ;; PRIMARY: (func $array.new_elem + ;; PRIMARY-NEXT: (drop + ;; PRIMARY-NEXT: (array.new_elem $elem-array $elem + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: (i32.const 0) + ;; PRIMARY-NEXT: ) + ;; PRIMARY-NEXT: ) + ;; PRIMARY-NEXT: ) + (func $array.new_elem + (drop + (array.new_elem $elem-array $elem + (i32.const 0) + (i32.const 0) + ) + ) + ) + + ;; SECONDARY: (type $0 (func)) + + ;; SECONDARY: (import "primary" "memory" (memory $mem 0)) + + ;; SECONDARY: (func $no-segment + ;; SECONDARY-NEXT: (nop) + ;; SECONDARY-NEXT: ) + (func $no-segment + (nop) + ) +) |