[wasm] Do not traverse all loop nodes in loop unrolling

The first pass in loop-unrolling aims to process stack check and loop
exit nodes. These nodes are all connected in a known fashion to the loop
header, so there is no need to traverse the whole loop.

Bug: v8:11298
Change-Id: I8336290ce7dddc8967f3fe05d3064b681c1e5c5a
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2732007Reviewed-by: Nico Hartmann <nicohartmann@chromium.org>
Commit-Queue: Manos Koukoutos <manoskouk@chromium.org>
Cr-Commit-Position: refs/heads/master@{#73166}

src/compiler/loop-unrolling.cc

@@ -44,18 +44,21 @@ void UnrollLoop(Node* loop_node, ZoneUnorderedSet<Node*>* loop, uint32_t depth,
 #define COPY(node, n) copier.map(node, n)
 #define FOREACH_COPY_INDEX(i) for (uint32_t i = 0; i < unrolling_count; i++)
 
-  for (Node* node : *loop) {
+  for (Node* node : loop_node->uses()) {
     switch (node->opcode()) {
-      case IrOpcode::kStackPointerGreaterThan: {
+      case IrOpcode::kBranch: {
         /*** Step 1: Remove stack checks from all but the first iteration of the
              loop. ***/
-        for (Edge edge : node->use_edges()) {
-          if (edge.from()->opcode() == IrOpcode::kBranch) {
-            FOREACH_COPY_INDEX(i) {
-              COPY(edge.from(), i)
-                  ->ReplaceInput(0, graph->NewNode(common->Int32Constant(1)));
-            }
-          } else if (edge.from()->opcode() == IrOpcode::kEffectPhi) {
+        Node* stack_check = node->InputAt(0);
+        if (stack_check->opcode() != IrOpcode::kStackPointerGreaterThan) {
+          break;
+        }
+        FOREACH_COPY_INDEX(i) {
+          COPY(node, i)->ReplaceInput(0,
+                                      graph->NewNode(common->Int32Constant(1)));
+        }
+        for (Node* use : stack_check->uses()) {
+          if (use->opcode() == IrOpcode::kEffectPhi) {
             // We now need to remove stack check and the related function call
             // from the effect chain.
             // The effect chain looks like this (* stand for irrelevant nodes):
@@ -68,24 +71,22 @@ void UnrollLoop(Node* loop_node, ZoneUnorderedSet<Node*>* loop, uint32_t depth,
             // | | | |
             // ( Load )
             //   |  |
-            // stack check
+            // stack_check
             // |  * | *
             // |  | | |
             // |  (call)
             // |   |     *
             // |   |     |
-            // stack check effect (that we need to replace)
-            Node* stack_check_effect = edge.from();
-            DCHECK_EQ(edge.index(), 0);
-            DCHECK_EQ(stack_check_effect->InputAt(1)->opcode(),
-                      IrOpcode::kCall);
-            DCHECK_EQ(stack_check_effect->InputAt(1)->InputAt(1), node);
-            DCHECK_EQ(node->InputAt(1)->opcode(), IrOpcode::kLoad);
-            DCHECK_EQ(node->InputAt(1)->InputAt(2)->opcode(), IrOpcode::kLoad);
-            Node* replacing_effect = node->InputAt(1)->InputAt(2)->InputAt(2);
+            // {use} (stack check effect that we need to replace)
+            DCHECK_EQ(use->InputAt(1)->opcode(), IrOpcode::kCall);
+            DCHECK_EQ(use->InputAt(1)->InputAt(1), stack_check);
+            DCHECK_EQ(stack_check->InputAt(1)->opcode(), IrOpcode::kLoad);
+            DCHECK_EQ(stack_check->InputAt(1)->InputAt(2)->opcode(),
+                      IrOpcode::kLoad);
+            Node* replacing_effect =
+                stack_check->InputAt(1)->InputAt(2)->InputAt(2);
             FOREACH_COPY_INDEX(i) {
-              COPY(stack_check_effect, i)
-                  ->ReplaceUses(COPY(replacing_effect, i));
+              COPY(use, i)->ReplaceUses(COPY(replacing_effect, i));
             }
           }
         }