[LICM]: support hosting ref_element_addr even if they are not guaranteed to be executed

In some instances, some instructions, like ref_element_addr, can be hoisted outside of loops even if they are not guaranteed to be executed.

We currently don’t support that / bail. We only try to do so / do further analysis only for begin_access because they are extremely heavy.

However, we need to support hosting of ref_element_addr in that case, if it does not have a loop dependent operand, in order to be able to hoist begin_access instructions in some benchmarks.

Initial local testing shows that this PR, when we enable exclusivity, improves the performance of a certain internal benchmark by over 40%

See rdar://problem/43623829

Author: Joe Shajrawi

lib/SILOptimizer/LoopTransforms/LICM.cpp

@@ -332,20 +332,23 @@ hoistSpecialInstruction(std::unique_ptr<LoopNestSummary> &LoopSummary,
   bool Changed = false;
 
   for (auto *Inst : Special) {
-    auto *BI = dyn_cast<BeginAccessInst>(Inst);
-    assert(BI && "Only BeginAccessInst are supported");
-    SmallVector<EndAccessInst *, 2> Ends;
-    getEndAccesses(BI, Ends);
-    if (!hoistInstruction(DT, BI, Loop, Preheader)) {
+    if (!hoistInstruction(DT, Inst, Loop, Preheader)) {
       continue;
     }
-    LLVM_DEBUG(llvm::dbgs() << "Hoisted " << *BI);
-    for (auto *instSink : Ends) {
-      if (!sinkInstruction(DT, LoopSummary, instSink, LI)) {
-        llvm_unreachable("LICM: Could not perform must-sink instruction");
+    if (auto *BI = dyn_cast<BeginAccessInst>(Inst)) {
+      SmallVector<EndAccessInst *, 2> Ends;
+      getEndAccesses(BI, Ends);
+      LLVM_DEBUG(llvm::dbgs() << "Hoisted BeginAccess " << *BI);
+      for (auto *instSink : Ends) {
+        if (!sinkInstruction(DT, LoopSummary, instSink, LI)) {
+          llvm_unreachable("LICM: Could not perform must-sink instruction");
+        }
       }
+      LLVM_DEBUG(llvm::errs() << " Successfully hosited and sank pair\n");
+    } else {
+      auto *REA = static_cast<RefElementAddrInst *>(Inst);
+      LLVM_DEBUG(llvm::dbgs() << "Hoisted RefElementAddr " << *REA);
     }
-    LLVM_DEBUG(llvm::errs() << " Successfully hosited and sank pair\n");
     Changed = true;
   }
 
@@ -615,6 +618,11 @@ void LoopTreeOptimization::analyzeCurrentLoop(
         checkSideEffects(Inst, MayWrites);
         break;
       }
+      case SILInstructionKind::RefElementAddrInst: {
+        auto *REA = static_cast<RefElementAddrInst *>(&Inst);
+        SpecialHoist.insert(REA);
+        break;
+      }
       case swift::SILInstructionKind::CondFailInst: {
         // We can (and must) hoist cond_fail instructions if the operand is
         // invariant. We must hoist them so that we preserve memory safety. A

test/SILOptimizer/licm.sil

@@ -319,3 +319,35 @@ bb3:
   dealloc_stack %0 : $*Builtin.Int32
   return %9999 : $()
 }
+
+class RefElemClass {
+  var x : Int32
+
+  init()
+}
+
+// Check hoisting of ref_element_addr in conditional control flow (for exclusivity)
+//
+// CHECK-LABEL: sil @hoist_ref_elem
+// CHECK: bb0(%0 : $RefElemClass):
+// CHECK-NEXT:   ref_element_addr %0 : $RefElemClass, #RefElemClass.x
+// CHECK-NEXT:   br bb1
+sil @hoist_ref_elem : $@convention(thin) (RefElemClass) -> () {
+bb0(%0 : $RefElemClass):
+  br bb1
+
+// loop.
+bb1:
+  cond_br undef, bb2, bb3
+
+bb2:
+  cond_br undef, bb4, bb1
+
+bb3:
+  %x = ref_element_addr %0 : $RefElemClass, #RefElemClass.x
+  br bb1
+
+bb4:
+  %10 = tuple ()
+  return %10 : $()
+}