[wasm-gc] Path-based type tracking for wasm-gc nodes

This CL adds control-path type-tracking for wasm-gc nodes in the
WasmGCOperatorReducer. Nodes now use the types assigned to their
argument nodes, as well as the additional information tracked along
control paths.

Drive-by: Add support for multiple instances of the same node to
appear in control-path-state.

Bug: v8:7748
Change-Id: I73e8f84595609b3a5fb61a2bffeb973182d17676
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3717994Reviewed-by: Maya Lekova <mslekova@chromium.org>
Commit-Queue: Manos Koukoutos <manoskouk@chromium.org>
Cr-Commit-Position: refs/heads/main@{#81373}

src/compiler/control-path-state.h

@@ -37,7 +37,13 @@ class ControlPathState {
 
   // Returns the {NodeState} assigned to node, or the default value
   // {NodeState()} if it is not assigned.
-  NodeState LookupState(Node* node) const { return states_.Get(node); }
+  NodeState LookupState(Node* node) const {
+    for (size_t depth = blocks_.Size(); depth > 0; depth--) {
+      NodeState state = states_.Get({node, depth});
+      if (state.IsSet()) return state;
+    }
+    return {};
+  }
 
   // Adds a state in the current code block, or a new block if the block list is
   // empty.
@@ -60,16 +66,24 @@ class ControlPathState {
   }
 
  private:
+  using NodeWithPathDepth = std::pair<Node*, size_t>;
+
 #if DEBUG
   bool BlocksAndStatesInvariant() {
-    PersistentMap<Node*, NodeState> states_copy(states_);
+    PersistentMap<NodeWithPathDepth, NodeState> states_copy(states_);
+    size_t depth = blocks_.Size();
     for (auto block : blocks_) {
+      std::unordered_set<Node*> seen_this_block;
       for (NodeState state : block) {
         // Every element of blocks_ has to be in states_.
-        if (states_copy.Get(state.node) != state) return false;
-        states_copy.Set(state.node, {});
+        if (seen_this_block.count(state.node) == 0) {
+          if (states_copy.Get({state.node, depth}) != state) return false;
+          states_copy.Set({state.node, depth}, {});
+          seen_this_block.emplace(state.node);
         }
       }
+      depth--;
+    }
     // Every element of {states_} has to be in {blocks_}. We removed all
     // elements of blocks_ from states_copy, so if it is not empty, the
     // invariant fails.
@@ -82,7 +96,7 @@ class ControlPathState {
   // set of states. It should hold at any point that the contents of {blocks_}
   // and {states_} is the same, which is implemented in
   // {BlocksAndStatesInvariant}.
-  PersistentMap<Node*, NodeState> states_;
+  PersistentMap<NodeWithPathDepth, NodeState> states_;
 };
 
 template <typename NodeState>
@@ -125,8 +139,6 @@ template <typename NodeState>
 void ControlPathState<NodeState>::AddState(Zone* zone, Node* node,
                                            NodeState state,
                                            ControlPathState<NodeState> hint) {
-  if (LookupState(node).IsSet()) return;
-
   FunctionalList<NodeState> prev_front = blocks_.Front();
   if (hint.blocks_.Size() > 0) {
     prev_front.PushFront(state, zone, hint.blocks_.Front());
@@ -135,7 +147,7 @@ void ControlPathState<NodeState>::AddState(Zone* zone, Node* node,
   }
   blocks_.DropFront();
   blocks_.PushFront(prev_front, zone);
-  states_.Set(node, state);
+  states_.Set({node, blocks_.Size()}, state);
   SLOW_DCHECK(BlocksAndStatesInvariant());
 }
 
@@ -143,10 +155,8 @@ template <typename NodeState>
 void ControlPathState<NodeState>::AddStateInNewBlock(Zone* zone, Node* node,
                                                      NodeState state) {
   FunctionalList<NodeState> new_block;
-  if (!LookupState(node).IsSet()) {
   new_block.PushFront(state, zone);
-    states_.Set(node, state);
-  }
+  states_.Set({node, blocks_.Size() + 1}, state);
   blocks_.PushFront(new_block, zone);
   SLOW_DCHECK(BlocksAndStatesInvariant());
 }
@@ -157,13 +167,13 @@ void ControlPathState<NodeState>::ResetToCommonAncestor(
   while (other.blocks_.Size() > blocks_.Size()) other.blocks_.DropFront();
   while (blocks_.Size() > other.blocks_.Size()) {
     for (NodeState state : blocks_.Front()) {
-      states_.Set(state.node, {});
+      states_.Set({state.node, blocks_.Size()}, {});
     }
     blocks_.DropFront();
   }
   while (blocks_ != other.blocks_) {
     for (NodeState state : blocks_.Front()) {
-      states_.Set(state.node, {});
+      states_.Set({state.node, blocks_.Size()}, {});
     }
     blocks_.DropFront();
     other.blocks_.DropFront();

src/compiler/pipeline.cc

@@ -2081,7 +2081,8 @@ struct WasmGCOptimizationPhase {
     GraphReducer graph_reducer(
         temp_zone, data->graph(), &data->info()->tick_counter(), data->broker(),
         data->jsgraph()->Dead(), data->observe_node_manager());
-    WasmGCOperatorReducer wasm_gc(&graph_reducer, data->mcgraph(), module);
+    WasmGCOperatorReducer wasm_gc(&graph_reducer, temp_zone, data->mcgraph(),
+                                  module);
     AddReducer(data, &graph_reducer, &wasm_gc);
     graph_reducer.ReduceGraph();
   }

src/compiler/simplified-operator.cc

@@ -1154,16 +1154,18 @@ struct SimplifiedOperatorGlobalCache final {
   LoadStackArgumentOperator kLoadStackArgument;
 
 #if V8_ENABLE_WEBASSEMBLY
+  // Note: The following two operators have a control input solely to find the
+  // typing context from the control path in wasm-gc-operator-reducer.
   struct IsNullOperator final : public Operator {
     IsNullOperator()
-        : Operator(IrOpcode::kIsNull, Operator::kPure, "IsNull", 1, 0, 0, 1, 0,
+        : Operator(IrOpcode::kIsNull, Operator::kPure, "IsNull", 1, 0, 1, 1, 0,
                    0) {}
   };
   IsNullOperator kIsNull;
 
   struct IsNotNullOperator final : public Operator {
     IsNotNullOperator()
-        : Operator(IrOpcode::kIsNotNull, Operator::kPure, "IsNotNull", 1, 0, 0,
+        : Operator(IrOpcode::kIsNotNull, Operator::kPure, "IsNotNull", 1, 0, 1,
                    1, 0, 0) {}
   };
   IsNotNullOperator kIsNotNull;

src/compiler/wasm-gc-operator-reducer.h

@@ -9,8 +9,11 @@
 #ifndef V8_COMPILER_WASM_GC_OPERATOR_REDUCER_H_
 #define V8_COMPILER_WASM_GC_OPERATOR_REDUCER_H_
 
+#include "src/compiler/control-path-state.h"
 #include "src/compiler/graph-reducer.h"
+#include "src/compiler/persistent-map.h"
 #include "src/compiler/wasm-graph-assembler.h"
+#include "src/wasm/wasm-subtyping.h"
 
 namespace v8 {
 namespace internal {
@@ -18,9 +21,28 @@ namespace compiler {
 
 class MachineGraph;
 
-class WasmGCOperatorReducer final : public AdvancedReducer {
+struct NodeWithType {
+  NodeWithType() : node(nullptr), type(wasm::kWasmVoid, nullptr) {}
+  NodeWithType(Node* node, wasm::TypeInModule type) : node(node), type(type) {}
+
+  bool operator==(const NodeWithType& other) const {
+    return node == other.node && type == other.type;
+  }
+  bool operator!=(const NodeWithType& other) const { return !(*this == other); }
+
+  bool IsSet() { return node != nullptr; }
+
+  Node* node;
+  wasm::TypeInModule type;
+};
+
+// This class optimizes away wasm-gc nodes based on the types of their
+// arguments. Although types have been assigned to nodes already, this class
+// also tracks additional type information along control paths.
+class WasmGCOperatorReducer final
+    : public AdvancedReducerWithControlPathState<NodeWithType> {
  public:
-  WasmGCOperatorReducer(Editor* editor, MachineGraph* mcgraph,
+  WasmGCOperatorReducer(Editor* editor, Zone* temp_zone_, MachineGraph* mcgraph,
                         const wasm::WasmModule* module);
 
   const char* reducer_name() const override { return "WasmGCOperatorReducer"; }
@@ -28,12 +50,22 @@ class WasmGCOperatorReducer final : public AdvancedReducer {
   Reduction Reduce(Node* node) final;
 
  private:
+  using ControlPathTypes = ControlPathState<NodeWithType>;
+
   Reduction ReduceAssertNotNull(Node* node);
-  Reduction ReduceIsNull(Node* node);
+  Reduction ReduceCheckNull(Node* node);
   Reduction ReduceWasmTypeCheck(Node* node);
   Reduction ReduceWasmTypeCast(Node* node);
+  Reduction ReduceMerge(Node* node);
+  Reduction ReduceIf(Node* node, bool condition);
+  Reduction ReduceStart(Node* node);
 
   Node* SetType(Node* node, wasm::ValueType type);
+  wasm::TypeInModule ObjectTypeFromContext(Node* object, Node* control);
+  Reduction UpdateNodeAndAliasesTypes(Node* state_owner,
+                                      ControlPathTypes parent_state, Node* node,
+                                      wasm::TypeInModule type,
+                                      bool in_new_block);
 
   Graph* graph() { return mcgraph_->graph(); }
   CommonOperatorBuilder* common() { return mcgraph_->common(); }

src/compiler/wasm-graph-assembler.cc

@@ -359,11 +359,11 @@ Node* WasmGraphAssembler::Null() {
 }
 
 Node* WasmGraphAssembler::IsNull(Node* object) {
-  return AddNode(graph()->NewNode(simplified_.IsNull(), object));
+  return AddNode(graph()->NewNode(simplified_.IsNull(), object, control()));
 }
 
 Node* WasmGraphAssembler::IsNotNull(Node* object) {
-  return AddNode(graph()->NewNode(simplified_.IsNotNull(), object));
+  return AddNode(graph()->NewNode(simplified_.IsNotNull(), object, control()));
 }
 
 Node* WasmGraphAssembler::AssertNotNull(Node* object) {

src/compiler/wasm-typer.cc

@@ -112,8 +112,8 @@ Reduction WasmTyper::Reduce(Node* node) {
         // AssertNotNull: Reverse the order of these operations, as this will
         // unlock more optimizations later.
         // We are implementing this in the typer so we can retype the nodes.
-        if (control->opcode() == IrOpcode::kWasmTypeCast && effect == object &&
-            control == object &&
+        while (control->opcode() == IrOpcode::kWasmTypeCast &&
+               effect == object && control == object &&
                !NodeProperties::GetType(object).AsWasm().type.is_bottom()) {
           Node* initial_object = NodeProperties::GetValueInput(object, 0);
           Node* previous_control = NodeProperties::GetControlInput(object);
@@ -128,7 +128,11 @@ Reduction WasmTyper::Reduce(Node* node) {
           object->ReplaceInput(NodeProperties::FirstValueIndex(object), node);
           object->ReplaceInput(NodeProperties::FirstEffectIndex(object), node);
           object->ReplaceInput(NodeProperties::FirstControlIndex(object), node);
+          Revisit(node);
           Revisit(object);
+          object = initial_object;
+          control = previous_control;
+          effect = previous_effect;
         }
       }
 

test/mjsunit/wasm/gc-optimizations.js

@@ -403,3 +403,68 @@ d8.file.execute("test/mjsunit/wasm/wasm-module-builder.js");
   let instance = builder.instantiate();
   assertEquals(10, instance.exports.main(10));
 })();
+
+(function PathBasedTypedOptimization() {
+  print(arguments.callee.name);
+  var builder = new WasmModuleBuilder();
+
+  let super_struct = builder.addStruct([makeField(kWasmI32, true)]);
+  let mid_struct = builder.addStruct(
+    [makeField(kWasmI32, true), makeField(kWasmI32, true)], super_struct);
+  let sub_struct = builder.addStruct(
+    [makeField(kWasmI32, true), makeField(kWasmI32, true),
+     makeField(kWasmI32, true)],
+    mid_struct);
+
+  let addToLocal = [kExprLocalGet, 1, kExprI32Add, kExprLocalSet, 1];
+
+  builder.addFunction(
+      "main", makeSig([wasmOptRefType(super_struct)], [kWasmI32]))
+    .addLocals(kWasmI32, 1)
+    .addBody([
+      kExprLocalGet, 0,
+      kGCPrefix, kExprRefTestStatic, sub_struct,
+
+      // These casts have to be preserved.
+      kExprLocalGet, 0,
+      kGCPrefix, kExprRefCastStatic, mid_struct,
+      kGCPrefix, kExprRefCastStatic, sub_struct,
+      kGCPrefix, kExprStructGet, sub_struct, 1,
+      ...addToLocal,
+
+      kExprIf, kWasmVoid,
+        // Both these casts should be optimized away.
+        kExprLocalGet, 0,
+        kGCPrefix, kExprRefCastStatic, mid_struct,
+        kGCPrefix, kExprRefCastStatic, sub_struct,
+        kGCPrefix, kExprStructGet, sub_struct, 1,
+        ...addToLocal,
+
+        kExprBlock, kWasmOptRef, super_struct,
+          kExprLocalGet, 0,
+          // This should also get optimized away.
+          kGCPrefix, kExprBrOnCastStaticFail, 0, mid_struct,
+          // So should this, despite being represented by a TypeGuard alias.
+          kGCPrefix, kExprRefCastStatic, sub_struct,
+          kGCPrefix, kExprStructGet, sub_struct, 1,
+          ...addToLocal,
+          kExprLocalGet, 0,  // Due to the branch result type.
+        kExprEnd,
+        kExprDrop,
+      kExprElse,
+        // This (always trapping) cast should be preserved.
+        kExprLocalGet, 0,
+        kGCPrefix, kExprRefCastStatic, sub_struct,
+        kGCPrefix, kExprStructGet, sub_struct, 1,
+        ...addToLocal,
+      kExprEnd,
+      // This cast should be preserved.
+      kExprLocalGet, 0,
+      kGCPrefix, kExprRefCastStatic, sub_struct,
+      kGCPrefix, kExprStructGet, sub_struct, 1,
+      kExprLocalGet, 1, kExprI32Add
+    ])
+    .exportFunc();
+
+  builder.instantiate();
+})();