RemoveUnusedModuleElements: Optimize indirect calls when the table is immutable

src/ir/table-utils.cpp

@@ -79,4 +79,82 @@ bool usesExpressions(ElementSegment* curr, Module* module) {
   return !allElementsRefFunc || hasSpecializedType;
 }
 
+TableInfoMap computeTableInfo(Module& wasm, bool initialContentsImmutable) {
+  // Set up the initial info.
+  TableInfoMap tables;
+  if (wasm.tables.empty()) {
+    return tables;
+  }
+  for (auto& table : wasm.tables) {
+    tables[table->name].initialContentsImmutable = initialContentsImmutable;
+    tables[table->name].flatTable =
+      std::make_unique<TableUtils::FlatTable>(wasm, *table);
+  }
+
+  // Next, look at the imports and exports.
+
+  for (auto& table : wasm.tables) {
+    if (table->imported()) {
+      tables[table->name].mayBeModified = true;
+    }
+  }
+
+  for (auto& ex : wasm.exports) {
+    if (ex->kind == ExternalKind::Table) {
+      tables[*ex->getInternalName()].mayBeModified = true;
+    }
+  }
+
+  // Find which tables have sets, by scanning for instructions. Only do so if we
+  // might learn anything new.
+  auto hasUnmodifiableTable = false;
+  for (auto& [_, info] : tables) {
+    if (!info.mayBeModified) {
+      hasUnmodifiableTable = true;
+      break;
+    }
+  }
+  if (!hasUnmodifiableTable) {
+    return tables;
+  }
+
+  using TablesWithSet = std::unordered_set<Name>;
+
+  ModuleUtils::ParallelFunctionAnalysis<TablesWithSet> analysis(
+    wasm, [&](Function* func, TablesWithSet& tablesWithSet) {
+      if (func->imported()) {
+        return;
+      }
+
+      struct Finder : public PostWalker<Finder> {
+        TablesWithSet& tablesWithSet;
+
+        Finder(TablesWithSet& tablesWithSet) : tablesWithSet(tablesWithSet) {}
+
+        void visitTableSet(TableSet* curr) {
+          tablesWithSet.insert(curr->table);
+        }
+        void visitTableFill(TableFill* curr) {
+          tablesWithSet.insert(curr->table);
+        }
+        void visitTableCopy(TableCopy* curr) {
+          tablesWithSet.insert(curr->destTable);
+        }
+        void visitTableInit(TableInit* curr) {
+          tablesWithSet.insert(curr->table);
+        }
+      };
+
+      Finder(tablesWithSet).walkFunction(func);
+    });
+
+  for (auto& [_, names] : analysis.map) {
+    for (auto name : names) {
+      tables[name].mayBeModified = true;
+    }
+  }
+
+  return tables;
+}
+
 } // namespace wasm::TableUtils

src/ir/table-utils.h

@@ -120,6 +120,53 @@ std::set<Name> getFunctionsNeedingElemDeclare(Module& wasm);
 // do so, and some do not, depending on their type and use.)
 bool usesExpressions(ElementSegment* curr, Module* module);
 
+// Information about a table's optimizability.
+struct TableInfo {
+  // Whether the table may be modifed at runtime, either because it is imported
+  // or exported, or table.set operations exist for it in the code.
+  bool mayBeModified = false;
+
+  // Whether we can assume that the initial contents are immutable. That is, if
+  // a table looks like [a, b, c] in the wasm, and we see a call to index 1, we
+  // will assume it must call b. It is possible that the table is appended to,
+  // but in this mode we assume the initial contents are not overwritten. This
+  // is the case for output from LLVM, for example.
+  //
+  // This is a weaker property than mayBeModified (if the table cannot be
+  // modified at all, we can definitely assume the initial contents we see are
+  // not mutated), but is useful in the case that things are appended to the
+  // table (as e.g. dynamic linking does in Emscripten, which passes in a flag
+  // to set this mode; in general, this is an invariant about the program that
+  // we must be informed about, not one that we can infer - there can be
+  // table.sets, for example, and this property implies that those sets never
+  // overwrite initial data).
+  bool initialContentsImmutable = false;
+
+  std::unique_ptr<TableUtils::FlatTable> flatTable;
+
+  // Whether we can optimize using this table's data on the entry level, that
+  // is, individual entries in the table are known to us, so calls through the
+  // table with known indexes can be inferred, etc.
+  bool canOptimizeByEntry() const {
+    // To infer entries, we require:
+    //  * Either the table can't be modified at all, or it can be modified but
+    //    the initial contents are immutable (so we can optimize those
+    //    contents, even if other things might be appended later, which we
+    //    cannot infer).
+    //  * The table is flat (so we can see what is in it, by index).
+    return (!mayBeModified || initialContentsImmutable) && flatTable->valid;
+  }
+};
+
+// A map of tables to their info.
+using TableInfoMap = std::unordered_map<Name, TableInfo>;
+
+// Compute a map with table optimizability info. We can be told that the initial
+// contents of the tables are immutable (that is, existing data is not
+// overwritten, but new things may be appended).
+TableInfoMap computeTableInfo(Module& wasm,
+                              bool initialContentsImmutable = false);
+
 } // namespace wasm::TableUtils
 
 #endif // wasm_ir_table_h

src/passes/Directize.cpp

@@ -23,11 +23,8 @@
 //
 //   --pass-arg=directize-initial-contents-immutable
 //
-// then the initial tables' contents are assumed to be immutable. That is, if
-// a table looks like [a, b, c] in the wasm, and we see a call to index 1, we
-// will assume it must call b. It is possible that the table is appended to, but
-// in this mode we assume the initial contents are not overwritten. This is the
-// case for output from LLVM, for example.
+// then the initial tables' contents are assumed to be immutable (see
+// TableUtils::TableInfo).
 //
 
 #include <unordered_map>
@@ -46,40 +43,18 @@ namespace wasm {
 
 namespace {
 
-struct TableInfo {
-  // Whether the table may be modifed at runtime, either because it is imported
-  // or exported, or table.set operations exist for it in the code.
-  bool mayBeModified = false;
-
-  // Whether we can assume that the initial contents are immutable. See the
-  // toplevel comment.
-  bool initialContentsImmutable = false;
-
-  std::unique_ptr<TableUtils::FlatTable> flatTable;
-
-  bool canOptimize() const {
-    // We can optimize if:
-    //  * Either the table can't be modified at all, or it can be modified but
-    //    the initial contents are immutable (so we can optimize them).
-    //  * The table is flat.
-    return (!mayBeModified || initialContentsImmutable) && flatTable->valid;
-  }
-};
-
-using TableInfoMap = std::unordered_map<Name, TableInfo>;
-
 struct FunctionDirectizer : public WalkerPass<PostWalker<FunctionDirectizer>> {
   bool isFunctionParallel() override { return true; }
 
   std::unique_ptr<Pass> create() override {
     return std::make_unique<FunctionDirectizer>(tables);
   }
 
-  FunctionDirectizer(const TableInfoMap& tables) : tables(tables) {}
+  FunctionDirectizer(const TableUtils::TableInfoMap& tables) : tables(tables) {}
 
   void visitCallIndirect(CallIndirect* curr) {
     auto& table = tables.at(curr->table);
-    if (!table.canOptimize()) {
+    if (!table.canOptimizeByEntry()) {
       return;
     }
     // If the target is constant, we can emit a direct call.
@@ -114,7 +89,7 @@ struct FunctionDirectizer : public WalkerPass<PostWalker<FunctionDirectizer>> {
   }
 
 private:
-  const TableInfoMap& tables;
+  const TableUtils::TableInfoMap& tables;
 
   bool changedTypes = false;
 
@@ -123,7 +98,7 @@ struct FunctionDirectizer : public WalkerPass<PostWalker<FunctionDirectizer>> {
   // that is, whether we know a direct call target, or we know it will trap, or
   // if we know nothing.
   CallUtils::IndirectCallInfo getTargetInfo(Expression* target,
-                                            const TableInfo& table,
+                                            const TableUtils::TableInfo& table,
                                             CallIndirect* original) {
     auto* c = target->dynCast<Const>();
     if (!c) {
@@ -165,7 +140,7 @@ struct FunctionDirectizer : public WalkerPass<PostWalker<FunctionDirectizer>> {
   // with an unreachable.
   void makeDirectCall(const std::vector<Expression*>& operands,
                       Expression* c,
-                      const TableInfo& table,
+                      const TableUtils::TableInfo& table,
                       CallIndirect* original) {
     auto info = getTargetInfo(c, table, original);
     if (std::get_if<CallUtils::Unknown>(&info)) {
@@ -211,84 +186,18 @@ struct Directize : public Pass {
     auto initialContentsImmutable =
       hasArgument("directize-initial-contents-immutable");
 
-    // Set up the initial info.
-    TableInfoMap tables;
-    for (auto& table : module->tables) {
-      tables[table->name].initialContentsImmutable = initialContentsImmutable;
-      tables[table->name].flatTable =
-        std::make_unique<TableUtils::FlatTable>(*module, *table);
-    }
-
-    // Next, look at the imports and exports.
-
-    for (auto& table : module->tables) {
-      if (table->imported()) {
-        tables[table->name].mayBeModified = true;
-      }
-    }
-
-    for (auto& ex : module->exports) {
-      if (ex->kind == ExternalKind::Table) {
-        tables[*ex->getInternalName()].mayBeModified = true;
-      }
-    }
-
-    // This may already be enough information to know that we can't optimize
-    // anything. If so, skip scanning all the module contents.
-    auto canOptimize = [&]() {
-      for (auto& [_, info] : tables) {
-        if (info.canOptimize()) {
-          return true;
-        }
-      }
-      return false;
-    };
-
-    if (!canOptimize()) {
-      return;
-    }
-
-    // Find which tables have sets.
+    auto tables =
+      TableUtils::computeTableInfo(*module, initialContentsImmutable);
 
-    using TablesWithSet = std::unordered_set<Name>;
-
-    ModuleUtils::ParallelFunctionAnalysis<TablesWithSet> analysis(
-      *module, [&](Function* func, TablesWithSet& tablesWithSet) {
-        if (func->imported()) {
-          return;
-        }
-
-        struct Finder : public PostWalker<Finder> {
-          TablesWithSet& tablesWithSet;
-
-          Finder(TablesWithSet& tablesWithSet) : tablesWithSet(tablesWithSet) {}
-
-          void visitTableSet(TableSet* curr) {
-            tablesWithSet.insert(curr->table);
-          }
-          void visitTableFill(TableFill* curr) {
-            tablesWithSet.insert(curr->table);
-          }
-          void visitTableCopy(TableCopy* curr) {
-            tablesWithSet.insert(curr->destTable);
-          }
-          void visitTableInit(TableInit* curr) {
-            tablesWithSet.insert(curr->table);
-          }
-        };
-
-        Finder(tablesWithSet).walkFunction(func);
-      });
-
-    for (auto& [_, names] : analysis.map) {
-      for (auto name : names) {
-        tables[name].mayBeModified = true;
+    // Stop if we cannot optimize anything.
+    auto hasOptimizableTable = false;
+    for (auto& [_, info] : tables) {
+      if (info.canOptimizeByEntry()) {
+        hasOptimizableTable = true;
+        break;
       }
     }
-
-    // Perhaps the new information about tables with sets shows we cannot
-    // optimize.
-    if (!canOptimize()) {
+    if (!hasOptimizableTable) {
       return;
     }
 

src/passes/RemoveUnusedModuleElements.cpp

@@ -45,6 +45,7 @@
 #include "ir/module-utils.h"
 #include "ir/struct-utils.h"
 #include "ir/subtypes.h"
+#include "ir/table-utils.h"
 #include "ir/utils.h"
 #include "pass.h"
 #include "support/insert_ordered.h"
@@ -172,11 +173,6 @@ struct Noter : public PostWalker<Noter, UnifiedExpressionVisitor<Noter>> {
     // the heap type we call with.
     reference({ModuleElementKind::Table, curr->table});
     noteIndirectCall(curr->table, curr->heapType);
-    // Note a possible call of a function reference as well, as something might
-    // be written into the table during runtime. With precise tracking of what
-    // is written into the table we could do better here; we could also see
-    // which tables are immutable. TODO
-    noteCallRef(curr->heapType);
   }
 
   void visitCallRef(CallRef* curr) {
@@ -407,6 +403,8 @@ struct Analyzer {
 
   std::unordered_set<IndirectCall> usedIndirectCalls;
 
+  std::optional<TableUtils::TableInfoMap> tableInfoMap;
+
   void useIndirectCall(IndirectCall call) {
     auto [_, inserted] = usedIndirectCalls.insert(call);
     if (!inserted) {
@@ -422,6 +420,16 @@ struct Analyzer {
     for (auto& elem : flatTableInfoMap[table].typeElems[type]) {
       reference({ModuleElementKind::ElementSegment, elem});
     }
+
+    // Note a possible call of a function reference as well, if something else
+    // might be written into the table during runtime.
+    // TODO: Add an option for immutable initial content like Directize?
+    if (!tableInfoMap) {
+      tableInfoMap = TableUtils::computeTableInfo(*module);
+    }
+    if ((*tableInfoMap)[table].mayBeModified) {
+      useCallRefType(type);
+    }
   }
 
   void useRefFunc(Name func) {