8342946: Replace predicate walking code in Loop Unrolling with a predicate visitor

chhagedorn · chhagedorn · commit 5f338e9adbcf · 2024-11-11T06:18:25.000Z
Reviewed-by: roland, kvn
diff --git a/src/hotspot/share/opto/loopTransform.cpp b/src/hotspot/share/opto/loopTransform.cpp
@@ -1779,47 +1779,19 @@ bool IdealLoopTree::is_invariant(Node* n) const {
 
 // Search the Assertion Predicates added by loop predication and/or range check elimination and update them according
 // to the new stride.
-void PhaseIdealLoop::update_main_loop_assertion_predicates(Node* ctrl, CountedLoopNode* loop_head, Node* init,
-                                                           const int stride_con) {
-  Node* entry = ctrl;
-  Node* prev_proj = ctrl;
-  LoopNode* outer_loop_head = loop_head->skip_strip_mined();
-  IdealLoopTree* outer_loop = get_loop(outer_loop_head);
+void PhaseIdealLoop::update_main_loop_assertion_predicates(CountedLoopNode* main_loop_head) {
+  Node* init = main_loop_head->init_trip();
 
   // Compute the value of the loop induction variable at the end of the
   // first iteration of the unrolled loop: init + new_stride_con - init_inc
-  int new_stride_con = stride_con * 2;
-  Node* max_value = _igvn.intcon(new_stride_con);
-  set_ctrl(max_value, C->root());
-
-  while (entry != nullptr && entry->is_Proj() && entry->in(0)->is_If()) {
-    IfNode* iff = entry->in(0)->as_If();
-    ProjNode* proj = iff->proj_out(1 - entry->as_Proj()->_con);
-    if (!proj->unique_ctrl_out()->is_Halt()) {
-      break;
-    }
-    Node* bol = iff->in(1);
-    if (bol->is_OpaqueTemplateAssertionPredicate()) {
-      assert(assertion_predicate_has_loop_opaque_node(iff), "must find OpaqueLoop* nodes");
-      // This is a Template Assertion Predicate for the initial or last access.
-      // Create an Initialized Assertion Predicates for it accordingly:
-      // - For the initial access a[init] (same as before)
-      // - For the last access a[init+new_stride-orig_stride] (with the new unroll stride)
-      prev_proj = create_initialized_assertion_predicate(iff, init, max_value, prev_proj);
-    } else if (bol->is_OpaqueInitializedAssertionPredicate()) {
-      // This is one of the two Initialized Assertion Predicates:
-      // - For the initial access a[init]
-      // - For the last access a[init+old_stride-orig_stride]
-      // We could keep the one for the initial access but we do not know which one we currently have here. Just kill both.
-      _igvn.replace_input_of(iff, 1, _igvn.intcon(1));
-    }
-    assert(!bol->is_OpaqueNotNull() || !loop_head->is_main_loop(), "OpaqueNotNull should not be at main loop");
-    entry = entry->in(0)->in(0);
-  }
-  if (prev_proj != ctrl) {
-    _igvn.replace_input_of(outer_loop_head, LoopNode::EntryControl, prev_proj);
-    set_idom(outer_loop_head, prev_proj, dom_depth(outer_loop_head));
-  }
+  int unrolled_stride_con = main_loop_head->stride_con() * 2;
+  Node* unrolled_stride = _igvn.intcon(unrolled_stride_con);
+  set_ctrl(unrolled_stride, C->root());
+
+  Node* loop_entry = main_loop_head->skip_strip_mined()->in(LoopNode::EntryControl);
+  PredicateIterator predicate_iterator(loop_entry);
+  UpdateStrideForAssertionPredicates update_stride_for_assertion_predicates(unrolled_stride, this);
+  predicate_iterator.for_each(update_stride_for_assertion_predicates);
 }
 
 // Source Loop: Cloned   - peeled_loop_head
@@ -1936,7 +1908,7 @@ void PhaseIdealLoop::do_unroll(IdealLoopTree *loop, Node_List &old_new, bool adj
   assert(old_trip_count > 1 && (!adjust_min_trip || stride_p <=
     MIN2<int>(max_jint / 2 - 2, MAX2(1<<3, Matcher::max_vector_size(T_BYTE)) * loop_head->unrolled_count())), "sanity");
 
-  update_main_loop_assertion_predicates(ctrl, loop_head, init, stride_con);
+  update_main_loop_assertion_predicates(loop_head);
 
   // Adjust loop limit to keep valid iterations number after unroll.
   // Use (limit - stride) instead of (((limit - init)/stride) & (-2))*stride
diff --git a/src/hotspot/share/opto/loopnode.hpp b/src/hotspot/share/opto/loopnode.hpp
@@ -948,7 +948,7 @@ class PhaseIdealLoop : public PhaseTransform {
  private:
   DEBUG_ONLY(static void count_opaque_loop_nodes(Node* n, uint& init, uint& stride);)
   static void get_assertion_predicates(ParsePredicateSuccessProj* parse_predicate_proj, Unique_Node_List& list, bool get_opaque = false);
-  void update_main_loop_assertion_predicates(Node* ctrl, CountedLoopNode* loop_head, Node* init, int stride_con);
+  void update_main_loop_assertion_predicates(CountedLoopNode* main_loop_head);
   void initialize_assertion_predicates_for_peeled_loop(CountedLoopNode* peeled_loop_head,
                                                        CountedLoopNode* remaining_loop_head,
                                                        uint first_node_index_in_cloned_loop_body,
diff --git a/src/hotspot/share/opto/predicates.cpp b/src/hotspot/share/opto/predicates.cpp
@@ -152,7 +152,6 @@ void TemplateAssertionPredicate::rewire_loop_data_dependencies(IfTrueNode* targe
   }
 }
 
-
 // Template Assertion Predicates always have the dedicated OpaqueTemplateAssertionPredicate to identify them.
 bool TemplateAssertionPredicate::is_predicate(Node* node) {
   if (!may_be_assertion_predicate_if(node)) {
@@ -179,6 +178,24 @@ IfTrueNode* TemplateAssertionPredicate::clone_and_replace_init(Node* new_control
   return success_proj;
 }
 
+// Replace the input to OpaqueLoopStrideNode with 'new_stride' and leave the other nodes unchanged.
+void TemplateAssertionPredicate::replace_opaque_stride_input(Node* new_stride, PhaseIterGVN& igvn) const {
+  TemplateAssertionExpression expression(opaque_node());
+  expression.replace_opaque_stride_input(new_stride, igvn);
+}
+
+// Create a new Initialized Assertion Predicate from this template at 'new_control' and return the success projection
+// of the newly created Initialized Assertion Predicate.
+IfTrueNode* TemplateAssertionPredicate::initialize(PhaseIdealLoop* phase, Node* new_control) const {
+  assert(phase->assertion_predicate_has_loop_opaque_node(head()),
+         "must find OpaqueLoop* nodes for Template Assertion Predicate");
+  InitializedAssertionPredicateCreator initialized_assertion_predicate(phase);
+  IfTrueNode* success_proj = initialized_assertion_predicate.create_from_template(head(), new_control);
+  assert(!phase->assertion_predicate_has_loop_opaque_node(success_proj->in(0)->as_If()),
+         "Initialized Assertion Predicates do not have OpaqueLoop* nodes in the bool expression anymore");
+  return success_proj;
+}
+
 // Initialized Assertion Predicates always have the dedicated OpaqueInitiailizedAssertionPredicate node to identify
 // them.
 bool InitializedAssertionPredicate::is_predicate(Node* node) {
@@ -189,6 +206,12 @@ bool InitializedAssertionPredicate::is_predicate(Node* node) {
   return if_node->in(1)->is_OpaqueInitializedAssertionPredicate();
 }
 
+void InitializedAssertionPredicate::kill(PhaseIdealLoop* phase) const {
+  Node* true_con = phase->igvn().intcon(1);
+  phase->set_ctrl(true_con, phase->C->root());
+  phase->igvn().replace_input_of(_if_node, 1, true_con);
+}
+
 #ifdef ASSERT
 // Check that the block has at most one Parse Predicate and that we only find Regular Predicate nodes (i.e. IfProj,
 // If, or RangeCheck nodes).
@@ -388,6 +411,63 @@ TemplateAssertionExpression::clone(const TransformStrategyForOpaqueLoopNodes& tr
   return opaque_node_clone->as_OpaqueTemplateAssertionPredicate();
 }
 
+// This class is used to replace the input to OpaqueLoopStrideNode with a new node while leaving the other nodes
+// unchanged.
+class ReplaceOpaqueStrideInput : public StackObj {
+  PhaseIterGVN& _igvn;
+  Unique_Node_List _nodes_to_visit;
+
+ public:
+  ReplaceOpaqueStrideInput(OpaqueTemplateAssertionPredicateNode* start_node, PhaseIterGVN& igvn) : _igvn(igvn) {
+    _nodes_to_visit.push(start_node);
+  }
+  NONCOPYABLE(ReplaceOpaqueStrideInput);
+
+  void replace(Node* new_opaque_stride_input) {
+    for (uint i = 0; i < _nodes_to_visit.size(); i++) {
+      Node* next = _nodes_to_visit[i];
+      for (uint j = 1; j < next->req(); j++) {
+        Node* input = next->in(j);
+        if (input->is_OpaqueLoopStride()) {
+          assert(TemplateAssertionExpressionNode::is_maybe_in_expression(input), "must also pass node filter");
+          _igvn.replace_input_of(input, 1, new_opaque_stride_input);
+        } else if (TemplateAssertionExpressionNode::is_maybe_in_expression(input)) {
+          _nodes_to_visit.push(input);
+        }
+      }
+    }
+  }
+};
+
+// Replace the input to OpaqueLoopStrideNode with 'new_stride' and leave the other nodes unchanged.
+void TemplateAssertionExpression::replace_opaque_stride_input(Node* new_stride, PhaseIterGVN& igvn) {
+  ReplaceOpaqueStrideInput replace_opaque_stride_input(_opaque_node, igvn);
+  replace_opaque_stride_input.replace(new_stride);
+}
+
+// The transformations of this class fold the OpaqueLoop* nodes by returning their inputs.
+class RemoveOpaqueLoopNodesStrategy : public TransformStrategyForOpaqueLoopNodes {
+ public:
+  Node* transform_opaque_init(OpaqueLoopInitNode* opaque_init) const override {
+    return opaque_init->in(1);
+  }
+
+  Node* transform_opaque_stride(OpaqueLoopStrideNode* opaque_stride) const override {
+    return opaque_stride->in(1);
+  }
+};
+
+OpaqueInitializedAssertionPredicateNode*
+TemplateAssertionExpression::clone_and_fold_opaque_loop_nodes(Node* new_control, PhaseIdealLoop* phase) {
+  RemoveOpaqueLoopNodesStrategy remove_opaque_loop_nodes_strategy;
+  OpaqueTemplateAssertionPredicateNode* cloned_template_opaque = clone(remove_opaque_loop_nodes_strategy, new_control,
+                                                                       phase);
+  OpaqueInitializedAssertionPredicateNode* opaque_initialized_opaque =
+      new OpaqueInitializedAssertionPredicateNode(cloned_template_opaque->in(1)->as_Bool(), phase->C);
+  phase->register_new_node(opaque_initialized_opaque, new_control);
+  return opaque_initialized_opaque;
+}
+
 // Check if this node belongs a Template Assertion Expression (including OpaqueLoop* nodes).
 bool TemplateAssertionExpressionNode::is_in_expression(Node* node) {
   if (is_maybe_in_expression(node)) {
@@ -664,6 +744,19 @@ IfTrueNode* InitializedAssertionPredicateCreator::create_from_template(IfNode* t
                               NOT_PRODUCT(COMMA template_assertion_predicate->assertion_predicate_type()));
 }
 
+// Create a new Initialized Assertion Predicate from 'template_assertion_predicate' by cloning it but omitting the
+// OpaqueLoop*Notes (i.e. taking their inputs instead).
+IfTrueNode* InitializedAssertionPredicateCreator::create_from_template(IfNode* template_assertion_predicate,
+                                                                       Node* new_control) {
+  OpaqueTemplateAssertionPredicateNode* template_opaque =
+      template_assertion_predicate->in(1)->as_OpaqueTemplateAssertionPredicate();
+  TemplateAssertionExpression template_assertion_expression(template_opaque);
+  OpaqueInitializedAssertionPredicateNode* assertion_expression =
+      template_assertion_expression.clone_and_fold_opaque_loop_nodes(new_control, _phase);
+  return create_control_nodes(new_control, template_assertion_predicate->Opcode(), assertion_expression
+                              NOT_PRODUCT(COMMA template_assertion_predicate->assertion_predicate_type()));
+}
+
 // Create a new Initialized Assertion Predicate directly without a template.
 IfTrueNode* InitializedAssertionPredicateCreator::create(Node* operand, Node* new_control, const jint stride,
                                                          const int scale, Node* offset, Node* range NOT_PRODUCT(COMMA
@@ -768,7 +861,7 @@ void CreateAssertionPredicatesVisitor::visit(const TemplateAssertionPredicate& t
 
 // Create an Initialized Assertion Predicate from the provided Template Assertion Predicate.
 IfTrueNode* CreateAssertionPredicatesVisitor::initialize_from_template(
-const TemplateAssertionPredicate& template_assertion_predicate) const {
+    const TemplateAssertionPredicate& template_assertion_predicate) const {
   IfNode* template_head = template_assertion_predicate.head();
   IfTrueNode* initialized_predicate = _phase->create_initialized_assertion_predicate(template_head, _init, _stride,
                                                                                      _new_control);
@@ -783,3 +876,36 @@ IfTrueNode* CreateAssertionPredicatesVisitor::clone_template_and_replace_init_in
   _phase->register_new_node(opaque_init, _new_control);
   return template_assertion_predicate.clone_and_replace_init(_new_control, opaque_init, _phase);
 }
+
+// Clone the Template Assertion Predicate and set a new input for the OpaqueLoopStrideNode.
+void UpdateStrideForAssertionPredicates::visit(const TemplateAssertionPredicate& template_assertion_predicate) {
+  replace_opaque_stride_input(template_assertion_predicate);
+  Node* template_tail_control_out = template_assertion_predicate.tail()->unique_ctrl_out();
+  IfTrueNode* initialized_success_proj = initialize_from_updated_template(template_assertion_predicate);
+  connect_initialized_assertion_predicate(template_tail_control_out, initialized_success_proj);
+}
+
+// Replace the input to OpaqueLoopStrideNode with 'new_stride' and leave the other nodes unchanged.
+void UpdateStrideForAssertionPredicates::replace_opaque_stride_input(
+    const TemplateAssertionPredicate& template_assertion_predicate) const {
+  template_assertion_predicate.replace_opaque_stride_input(_new_stride, _phase->igvn());
+}
+
+IfTrueNode* UpdateStrideForAssertionPredicates::initialize_from_updated_template(
+    const TemplateAssertionPredicate& template_assertion_predicate) const {
+  IfTrueNode* initialized_success_proj = template_assertion_predicate.initialize(_phase, template_assertion_predicate.tail());
+  return initialized_success_proj;
+}
+
+// The newly created Initialized Assertion Predicate can safely be inserted because this visitor is already visiting
+// the Template Assertion Predicate above this. So, we will not accidentally visit this again and kill it with the
+// visit() method for Initialized Assertion Predicates.
+void UpdateStrideForAssertionPredicates::connect_initialized_assertion_predicate(
+    Node* new_control_out, IfTrueNode* initialized_success_proj) const {
+  if (new_control_out->is_Loop()) {
+    _phase->igvn().replace_input_of(new_control_out, LoopNode::EntryControl, initialized_success_proj);
+  } else {
+    _phase->igvn().replace_input_of(new_control_out, 0, initialized_success_proj);
+  }
+  _phase->set_idom(new_control_out, initialized_success_proj, _phase->dom_depth(new_control_out));
+}
diff --git a/src/hotspot/share/opto/predicates.hpp b/src/hotspot/share/opto/predicates.hpp
@@ -404,6 +404,8 @@ class TemplateAssertionPredicate : public Predicate {
   }
 
   IfTrueNode* clone_and_replace_init(Node* new_control, OpaqueLoopInitNode* new_opaque_init, PhaseIdealLoop* phase) const;
+  void replace_opaque_stride_input(Node* new_stride, PhaseIterGVN& igvn) const;
+  IfTrueNode* initialize(PhaseIdealLoop* phase, Node* new_control) const;
   void rewire_loop_data_dependencies(IfTrueNode* target_predicate, const NodeInLoopBody& data_in_loop_body,
                                      PhaseIdealLoop* phase) const;
   static bool is_predicate(Node* node);
@@ -434,6 +436,7 @@ class InitializedAssertionPredicate : public Predicate {
     return _success_proj;
   }
 
+  void kill(PhaseIdealLoop* phase) const;
   static bool is_predicate(Node* node);
 };
 
@@ -461,6 +464,8 @@ class TemplateAssertionExpression : public StackObj {
   OpaqueTemplateAssertionPredicateNode* clone_and_replace_init(Node* new_init, Node* new_ctrl, PhaseIdealLoop* phase);
   OpaqueTemplateAssertionPredicateNode* clone_and_replace_init_and_stride(Node* new_control, Node* new_init,
                                                                           Node* new_stride, PhaseIdealLoop* phase);
+  void replace_opaque_stride_input(Node* new_stride, PhaseIterGVN& igvn);
+  OpaqueInitializedAssertionPredicateNode* clone_and_fold_opaque_loop_nodes(Node* new_ctrl, PhaseIdealLoop* phase);
 };
 
 // Class to represent a node being part of a Template Assertion Expression. Note that this is not an IR node.
@@ -608,6 +613,7 @@ class InitializedAssertionPredicateCreator : public StackObj {
 
   IfTrueNode* create_from_template(IfNode* template_assertion_predicate, Node* new_control, Node* new_init,
                                    Node* new_stride);
+  IfTrueNode* create_from_template(IfNode* template_assertion_predicate, Node* new_control);
   IfTrueNode* create(Node* operand, Node* new_control, jint stride, int scale, Node* offset, Node* range
                      NOT_PRODUCT(COMMA AssertionPredicateType assertion_predicate_type));
 
@@ -1032,4 +1038,32 @@ class TemplateAssertionPredicateCollector : public PredicateVisitor {
   }
 };
 
+// This visitor updates the stride for an Assertion Predicate during Loop Unrolling. The inputs to the OpaqueLoopStride
+// nodes Template of Template Assertion Predicates are updated and new Initialized Assertion Predicates are created
+// from the updated templates. The old Initialized Assertion Predicates are killed.
+class UpdateStrideForAssertionPredicates : public PredicateVisitor {
+  Node* const _new_stride;
+  PhaseIdealLoop* const _phase;
+
+  void replace_opaque_stride_input(const TemplateAssertionPredicate& template_assertion_predicate) const;
+  IfTrueNode* initialize_from_updated_template(const TemplateAssertionPredicate& template_assertion_predicate) const;
+  void connect_initialized_assertion_predicate(Node* new_control_out, IfTrueNode* initialized_success_proj) const;
+
+ public:
+  UpdateStrideForAssertionPredicates(Node* const new_stride, PhaseIdealLoop* phase)
+      : _new_stride(new_stride),
+        _phase(phase) {}
+  NONCOPYABLE(UpdateStrideForAssertionPredicates);
+
+  using PredicateVisitor::visit;
+
+  void visit(const TemplateAssertionPredicate& template_assertion_predicate) override;
+
+  // Kill the old Initialized Assertion Predicates with old strides before unrolling. The new Initialized Assertion
+  // Predicates are inserted after the Template Assertion Predicate which ensures that we are not accidentally visiting
+  // and killing a newly created Initialized Assertion Predicate here.
+  void visit(const InitializedAssertionPredicate& initialized_assertion_predicate) override {
+    initialized_assertion_predicate.kill(_phase);
+  }
+};
 #endif // SHARE_OPTO_PREDICATES_HPP
