diff --git a/src/ir/module-splitting.cpp b/src/ir/module-splitting.cpp index 21695f0bf07..0d11829e034 100644 --- a/src/ir/module-splitting.cpp +++ b/src/ir/module-splitting.cpp @@ -77,6 +77,7 @@ #include "ir/find_all.h" #include "ir/module-utils.h" #include "ir/names.h" +#include "support/small_vector.h" #include "support/stdckdint.h" #include "wasm-builder.h" #include "wasm.h" @@ -313,6 +314,161 @@ TableSlotManager::Slot TableSlotManager::getSlot(Name func, HeapType type) { return newSlot; } +// Module items ownership tracking + +// Struct containing sets of used module elements of a single module +struct UsedNames { + std::unordered_set globals; + std::unordered_set memories; + std::unordered_set tables; + std::unordered_set tags; + std::unordered_set dataSegments; + std::unordered_set elementSegments; +}; + +// A tracker that, given a module element, tracks which module is its owner, +// i.e., where the element should be placed, and the list of secondary modules +// using this element. +struct OwnershipTracker { + UsedNames primaryUsed; + std::vector secondaryUsed; + + struct ItemInfo { + UsedNames* owner = nullptr; + SmallVector usingSecondaries; + }; + + std::unordered_map tables; + std::unordered_map memories; + std::unordered_map globals; + std::unordered_map tags; + std::unordered_map dataSegments; + std::unordered_map elementSegments; + + const std::vector>* secondaries = nullptr; + + using FieldType = std::unordered_set UsedNames::*; + using MapType = std::unordered_map OwnershipTracker::*; + + template void insert(Name name, UsedNames* owner) { + if constexpr (std::is_same_v) { + insertImpl(name, owner, &OwnershipTracker::tables, &UsedNames::tables); + } else if constexpr (std::is_same_v) { + insertImpl( + name, owner, &OwnershipTracker::memories, &UsedNames::memories); + } else if constexpr (std::is_same_v) { + insertImpl(name, owner, &OwnershipTracker::globals, &UsedNames::globals); + } else if constexpr (std::is_same_v) { + insertImpl(name, owner, &OwnershipTracker::tags, &UsedNames::tags); + } else if constexpr (std::is_same_v) { + insertImpl( + name, owner, &OwnershipTracker::dataSegments, &UsedNames::dataSegments); + } else if constexpr (std::is_same_v) { + insertImpl(name, + owner, + &OwnershipTracker::elementSegments, + &UsedNames::elementSegments); + } + } + + // 'mapField' points to one of OwnershipTracker's maps, such as + // std::unordered_map globals; + // 'field' points to one of UsedName's sets, such as + // std::unordered_set globals; + void + insertImpl(Name name, UsedNames* owner, MapType mapField, FieldType field) { + (owner->*field).insert(name); + // Figure out which module is the 'owner' of this item. If it is used by a + // single secondary module, that secondary module is the owner. If it is + // used by the primary module or multiple secondary modules, the primary + // module is the owner. + auto [it, inserted] = (this->*mapField).insert({name, ItemInfo{owner, {}}}); + Module* secondary = nullptr; + if (owner != &primaryUsed) { + size_t index = owner - secondaryUsed.data(); + secondary = (*secondaries)[index].get(); + } + if (inserted) { + if (secondary) { + it->second.usingSecondaries.push_back(secondary); + } + } else { + if (it->second.owner != owner) { + it->second.owner = &primaryUsed; + (primaryUsed.*field).insert(name); + } + if (secondary) { + auto& vec = it->second.usingSecondaries; + if (std::find(vec.begin(), vec.end(), secondary) == vec.end()) { + vec.push_back(secondary); + } + } + } + } + + void build(const std::vector>& secondaries) { + this->secondaries = &secondaries; + + // Build initial maps of a module element Name to an ItemInfo for each + // module element type. + // 'field' points to one of UsedName's sets, such as + // std::unordered_set globals; + auto buildMap = [&](FieldType field, + std::unordered_map& map) { + for (auto& name : (primaryUsed.*field)) { + map[name].owner = &primaryUsed; + } + for (size_t i = 0; i < secondaryUsed.size(); ++i) { + auto& secUsed = secondaryUsed[i]; + auto* secondary = secondaries[i].get(); + for (auto& name : (secUsed.*field)) { + auto [it, inserted] = map.insert({name, ItemInfo{&secUsed, {}}}); + it->second.usingSecondaries.push_back(secondary); + if (!inserted) { + it->second.owner = &primaryUsed; + } + } + } + }; + buildMap(&UsedNames::tables, tables); + buildMap(&UsedNames::memories, memories); + buildMap(&UsedNames::globals, globals); + buildMap(&UsedNames::tags, tags); + buildMap(&UsedNames::dataSegments, dataSegments); + buildMap(&UsedNames::elementSegments, elementSegments); + } + + UsedNames* getOwner(Name name, + const std::unordered_map& map) { + auto it = map.find(name); + if (it != map.end()) { + return it->second.owner; + } + return nullptr; + } + + const SmallVector& + getUsingSecondaries(Name name, + const std::unordered_map& map) { + auto it = map.find(name); + if (it != map.end()) { + return it->second.usingSecondaries; + } + static SmallVector empty; + return empty; + } + + bool useEmpty(Name name, const std::unordered_map& map) { + return getOwner(name, map) == nullptr; + } + + bool usedBySingleSecondary(Name name, + const std::unordered_map& map) { + auto* owner = getOwner(name, map); + return owner != nullptr && owner != &primaryUsed; + } +}; + struct ModuleSplitter { const Config& config; std::vector> secondaries; @@ -359,17 +515,9 @@ struct ModuleSplitter { ExternalKind kind); Name getTrampoline(Name funcName); - struct UsedNames { - std::unordered_set globals; - std::unordered_set memories; - std::unordered_set tables; - std::unordered_set tags; - std::unordered_set dataSegments; - std::unordered_set elementSegments; - }; - using PrimarySecondaryUsedNames = - std::pair>; - PrimarySecondaryUsedNames computeUsedNames(); + OwnershipTracker tracker; + + void computeUsedNames(); // Main splitting steps void classifyFunctions(); @@ -649,12 +797,20 @@ void ModuleSplitter::thunkExportedSecondaryFunctions() { } } -ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { +void ModuleSplitter::computeUsedNames() { + tracker.secondaries = &secondaries; + tracker.secondaryUsed.resize(secondaries.size()); + UsedNames& primaryUsed = tracker.primaryUsed; + std::vector& secondaryUsed = tracker.secondaryUsed; + struct NameCollector : public PostWalker> { UsedNames& used; - NameCollector(UsedNames& used) : used(used) {} + OwnershipTracker& tracker; + + NameCollector(UsedNames& used, OwnershipTracker& tracker) + : used(used), tracker(tracker) {} void visitExpression(Expression* curr) { #define DELEGATE_ID curr->_id @@ -674,22 +830,22 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (cast->field.is()) { \ switch (kind) { \ case ModuleItemKind::Table: \ - used.tables.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::Memory: \ - used.memories.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::Global: \ - used.globals.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::Tag: \ - used.tags.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::DataSegment: \ - used.dataSegments.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::ElementSegment: \ - used.elementSegments.insert(cast->field); \ + tracker.insert(cast->field, &used); \ break; \ case ModuleItemKind::Function: \ case ModuleItemKind::Invalid: \ @@ -702,21 +858,18 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { }; // Given a module, collect names used in the module - auto scanModule = [&](Module& module) { - UsedNames used; - NameCollector collector(used); + auto scanModule = [&](Module& module, UsedNames& used) { + NameCollector collector(used, tracker); for (auto& func : module.functions) { if (!func->imported()) { collector.walk(func->body); } } - return used; }; - UsedNames primaryUsed = scanModule(primary); - std::vector secondaryUsed; - for (auto& secondaryPtr : secondaries) { - secondaryUsed.push_back(scanModule(*secondaryPtr)); + scanModule(primary, primaryUsed); + for (size_t i = 0; i < secondaries.size(); ++i) { + scanModule(*secondaries[i], secondaryUsed[i]); } // If primary module has exports, they are "used" in it. Secondary modules @@ -724,16 +877,16 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { for (auto& ex : primary.exports) { switch (ex->kind) { case ExternalKind::Global: - primaryUsed.globals.insert(*ex->getInternalName()); + tracker.insert(*ex->getInternalName(), &primaryUsed); break; case ExternalKind::Memory: - primaryUsed.memories.insert(*ex->getInternalName()); + tracker.insert(*ex->getInternalName(), &primaryUsed); break; case ExternalKind::Table: - primaryUsed.tables.insert(*ex->getInternalName()); + tracker.insert
(*ex->getInternalName(), &primaryUsed); break; case ExternalKind::Tag: - primaryUsed.tags.insert(*ex->getInternalName()); + tracker.insert(*ex->getInternalName(), &primaryUsed); break; default: break; @@ -743,10 +896,10 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { // We need to assume the dispatch table and its base global are used in the // primary module, because we will create segments there later. if (tableManager.dispatchTable) { - primaryUsed.tables.insert(tableManager.dispatchTable->name); + tracker.insert
(tableManager.dispatchTable->name, &primaryUsed); } if (tableManager.dispatchBase.global) { - primaryUsed.globals.insert(tableManager.dispatchBase.global); + tracker.insert(tableManager.dispatchBase.global, &primaryUsed); } // If custom-descirptors is enabled, global and table initializers can trap. @@ -757,40 +910,18 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (global->init && EffectAnalyzer(config.passOptions, primary, global->init) .hasUnremovableSideEffects()) { - primaryUsed.globals.insert(global->name); + tracker.insert(global->name, &primaryUsed); } } for (auto& table : primary.tables) { if (table->init && EffectAnalyzer(config.passOptions, primary, table->init) .hasUnremovableSideEffects()) { - primaryUsed.tables.insert(table->name); + tracker.insert
(table->name, &primaryUsed); } } } - // Given a name and a module item kind (field pointer), find which module - // "owns" it. If it is used by exactly one secondary module, that secondary - // module is the owner. If it is used by the primary module or multiple - // secondary modules, the primary module is the owner. If it is not used, - // returns nullptr. - auto getOwner = [&](Name name, auto UsedNames::* field) -> UsedNames* { - UsedNames* owner = nullptr; - if ((primaryUsed.*field).contains(name)) { - owner = &primaryUsed; - } - for (auto& sec : secondaryUsed) { - if ((sec.*field).contains(name)) { - if (owner) { - owner = &primaryUsed; - break; - } - owner = &sec; - } - } - return owner; - }; - // Scan table initializers into their owning modules. If a table is used by a // single secondary module, its initializer dependencies are marked as "used" // in that secondary module. Otherwise, they are marked as used in the primary @@ -800,8 +931,8 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (!table->init) { continue; } - if (UsedNames* owner = getOwner(table->name, &UsedNames::tables)) { - NameCollector(*owner).walk(table->init); + if (UsedNames* owner = tracker.getOwner(table->name, tracker.tables)) { + NameCollector(*owner, tracker).walk(table->init); } } } @@ -850,13 +981,16 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { return false; }; +#define ADD_ITEM_TO_TRACKER_TO_TRACKER(FIELD, VAL) \ + tracker.insert(VAL, owner, &OwnershipTracker::FIELD, &UsedNames::FIELD) + // Iterate on active data and element segments. If its table or memory is // used by a single secondary module, mark it "used" there. Only scan its // 'offset' or 'data'(in case of ElementSegment) and add it to that module's // used only when it is a sole secondary owner. If not assign it to the // primary module and scan it there. ModuleUtils::iterActiveDataSegments(primary, [&](DataSegment* segment) { - UsedNames* owner = getOwner(segment->memory, &UsedNames::memories); + UsedNames* owner = tracker.getOwner(segment->memory, tracker.memories); // Trapping segments should be kept in the primary module because they are // evaluated at the instantiation time. if (mayTrap(segment)) { @@ -865,15 +999,15 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (!owner) { return; } - owner->dataSegments.insert(segment->name); - owner->memories.insert(segment->memory); + tracker.insert(segment->name, owner); + tracker.insert(segment->memory, owner); if (segment->offset) { - NameCollector(*owner).walk(segment->offset); + NameCollector(*owner, tracker).walk(segment->offset); } }); ModuleUtils::iterActiveElementSegments(primary, [&](ElementSegment* segment) { - UsedNames* owner = getOwner(segment->table, &UsedNames::tables); + UsedNames* owner = tracker.getOwner(segment->table, tracker.tables); // If placeholders are NOT used, and if all functions in an element segment // belong to a single secondary module, we can move the segment to that @@ -920,13 +1054,13 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (!owner) { return; } - owner->elementSegments.insert(segment->name); - owner->tables.insert(segment->table); + tracker.insert(segment->name, owner); + tracker.insert
(segment->table, owner); if (segment->offset) { - NameCollector(*owner).walk(segment->offset); + NameCollector(*owner, tracker).walk(segment->offset); } for (auto* item : segment->data) { - NameCollector(*owner).walk(item); + NameCollector(*owner, tracker).walk(item); } }); @@ -938,7 +1072,7 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (segment->isPassive() && primaryUsed.elementSegments.contains(segment->name)) { for (auto* item : segment->data) { - NameCollector(primaryUsed).walk(item); + NameCollector(primaryUsed, tracker).walk(item); } } } @@ -955,32 +1089,16 @@ ModuleSplitter::PrimarySecondaryUsedNames ModuleSplitter::computeUsedNames() { if (!global->init) { continue; } - if (UsedNames* owner = getOwner(global->name, &UsedNames::globals)) { + if (UsedNames* owner = tracker.getOwner(global->name, tracker.globals)) { for (auto* get : FindAll(global->init).list) { - owner->globals.insert(get->name); + tracker.insert(get->name, owner); } } } - - return std::make_pair(primaryUsed, secondaryUsed); } void ModuleSplitter::shareImportableItems() { - auto usedNames = computeUsedNames(); - auto& primaryUsed = usedNames.first; - auto& secondaryUsed = usedNames.second; - - // Given a name and module item kind, returns the list of secondary modules - // using that name - auto getUsingSecondaries = [&](const Name& name, auto UsedNames::* field) { - std::vector usingModules; - for (size_t i = 0; i < secondaries.size(); ++i) { - if ((secondaryUsed[i].*field).contains(name)) { - usingModules.push_back(secondaries[i].get()); - } - } - return usingModules; - }; + computeUsedNames(); // Share module items with secondary modules. // 1. Only share an item with the modules that use it @@ -991,18 +1109,16 @@ void ModuleSplitter::shareImportableItems() { std::vector memoriesToRemove; for (auto& memory : primary.memories) { - auto usingSecondaries = - getUsingSecondaries(memory->name, &UsedNames::memories); - bool inPrimary = primaryUsed.memories.contains(memory->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(memory->name, tracker.memories)) { memoriesToRemove.push_back(memory->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(memory->name, tracker.memories)) { + auto* secondary = + tracker.getUsingSecondaries(memory->name, tracker.memories)[0]; ModuleUtils::copyMemory(memory.get(), *secondary); memoriesToRemove.push_back(memory->name); } else { - for (auto* secondary : usingSecondaries) { + for (auto* secondary : + tracker.getUsingSecondaries(memory->name, tracker.memories)) { auto* secondaryMemory = ModuleUtils::copyMemory(memory.get(), *secondary); makeImportExport( @@ -1016,19 +1132,17 @@ void ModuleSplitter::shareImportableItems() { std::vector tablesToRemove; for (auto& table : primary.tables) { - auto usingSecondaries = - getUsingSecondaries(table->name, &UsedNames::tables); - bool inPrimary = primaryUsed.tables.contains(table->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(table->name, tracker.tables)) { tablesToRemove.push_back(table->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(table->name, tracker.tables)) { + auto* secondary = + tracker.getUsingSecondaries(table->name, tracker.tables)[0]; assert(!secondary->getTableOrNull(table->name)); ModuleUtils::copyTable(table.get(), *secondary); tablesToRemove.push_back(table->name); } else { - for (auto* secondary : usingSecondaries) { + for (auto* secondary : + tracker.getUsingSecondaries(table->name, tracker.tables)) { auto* secondaryTable = ModuleUtils::copyTable(table.get(), *secondary); makeImportExport(*table, *secondaryTable, "table", ExternalKind::Table); } @@ -1045,18 +1159,16 @@ void ModuleSplitter::shareImportableItems() { "TODO: add wrapper functions for disallowed mutable globals"); } - auto usingSecondaries = - getUsingSecondaries(global->name, &UsedNames::globals); - bool inPrimary = primaryUsed.globals.contains(global->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(global->name, tracker.globals)) { globalsToRemove.push_back(global->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(global->name, tracker.globals)) { + auto* secondary = + tracker.getUsingSecondaries(global->name, tracker.globals)[0]; ModuleUtils::copyGlobal(global.get(), *secondary); globalsToRemove.push_back(global->name); } else { - for (auto* secondary : usingSecondaries) { + for (auto* secondary : + tracker.getUsingSecondaries(global->name, tracker.globals)) { auto* secondaryGlobal = ModuleUtils::copyGlobal(global.get(), *secondary); makeImportExport( @@ -1070,17 +1182,15 @@ void ModuleSplitter::shareImportableItems() { std::vector tagsToRemove; for (auto& tag : primary.tags) { - auto usingSecondaries = getUsingSecondaries(tag->name, &UsedNames::tags); - bool inPrimary = primaryUsed.tags.contains(tag->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(tag->name, tracker.tags)) { tagsToRemove.push_back(tag->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(tag->name, tracker.tags)) { + auto* secondary = tracker.getUsingSecondaries(tag->name, tracker.tags)[0]; ModuleUtils::copyTag(tag.get(), *secondary); tagsToRemove.push_back(tag->name); } else { - for (auto* secondary : usingSecondaries) { + for (auto* secondary : + tracker.getUsingSecondaries(tag->name, tracker.tags)) { auto* secondaryTag = ModuleUtils::copyTag(tag.get(), *secondary); makeImportExport(*tag, *secondaryTag, "tag", ExternalKind::Tag); } @@ -1096,14 +1206,12 @@ void ModuleSplitter::shareImportableItems() { std::vector dataSegmentsToRemove; for (auto& dataSegment : primary.dataSegments) { - auto usingSecondaries = - getUsingSecondaries(dataSegment->name, &UsedNames::dataSegments); - bool inPrimary = primaryUsed.dataSegments.contains(dataSegment->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(dataSegment->name, tracker.dataSegments)) { dataSegmentsToRemove.push_back(dataSegment->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(dataSegment->name, + tracker.dataSegments)) { + auto* secondary = + tracker.getUsingSecondaries(dataSegment->name, tracker.dataSegments)[0]; ModuleUtils::copyDataSegment(dataSegment.get(), *secondary); dataSegmentsToRemove.push_back(dataSegment->name); } @@ -1114,14 +1222,12 @@ void ModuleSplitter::shareImportableItems() { std::vector elementSegmentsToRemove; for (auto& elementSegment : primary.elementSegments) { - auto usingSecondaries = - getUsingSecondaries(elementSegment->name, &UsedNames::elementSegments); - bool inPrimary = primaryUsed.elementSegments.contains(elementSegment->name); - - if (!inPrimary && usingSecondaries.empty()) { + if (tracker.useEmpty(elementSegment->name, tracker.elementSegments)) { elementSegmentsToRemove.push_back(elementSegment->name); - } else if (!inPrimary && usingSecondaries.size() == 1) { - auto* secondary = usingSecondaries[0]; + } else if (tracker.usedBySingleSecondary(elementSegment->name, + tracker.elementSegments)) { + auto* secondary = tracker.getUsingSecondaries(elementSegment->name, + tracker.elementSegments)[0]; ModuleUtils::copyElementSegment(elementSegment.get(), *secondary); elementSegmentsToRemove.push_back(elementSegment->name); }