LoopUnroller: Also handle > exit condition

Author: aschwaighofer

lib/SILOptimizer/LoopTransforms/LoopUnroll.cpp

@@ -126,13 +126,23 @@ static Optional<uint64_t> getMaxLoopTripCount(SILLoop *Loop,
   IntegerLiteralInst *End;
   SILValue RecNext;
 
+  unsigned Adjust = 0;
+
   if (!match(CondBr->getCondition(),
              m_BuiltinInst(BuiltinValueKind::ICMP_EQ, m_SILValue(RecNext),
                            m_IntegerLiteralInst(End))) &&
       !match(CondBr->getCondition(),
              m_BuiltinInst(BuiltinValueKind::ICMP_SGE, m_SILValue(RecNext),
-                           m_IntegerLiteralInst(End))))
-    return None;
+                           m_IntegerLiteralInst(End)))) {
+    if (!match(CondBr->getCondition(),
+               m_BuiltinInst(BuiltinValueKind::ICMP_SGT, m_SILValue(RecNext),
+                             m_IntegerLiteralInst(End))))
+      return None;
+    // Otherwise, we have a greater than comparison.
+    else
+      Adjust = 1;
+  }
+
   if (!match(RecNext,
              m_TupleExtractInst(m_ApplyInst(BuiltinValueKind::SAddOver,
                                             m_SILPHIArgument(RecArg), m_One()),
@@ -162,7 +172,7 @@ static Optional<uint64_t> getMaxLoopTripCount(SILLoop *Loop,
   if (Dist == 0)
     return None;
 
-  return Dist.getZExtValue();
+  return Dist.getZExtValue() + Adjust;
 }
 
 /// Check whether we can duplicate the instructions in the loop and use a

test/SILOptimizer/loop_unroll.sil

@@ -107,6 +107,42 @@ bb3:
  return %8 : $()
 }
 
+// CHECK-LABEL: sil @loop_unroll_4
+// CHECK: bb0:
+// CHECK:  br bb1
+// CHECK: bb1
+// CHECK:  = builtin "sadd_with_overflow_Int64
+// CHECK:  cond_br {{.*}}, bb3, bb2
+// CHECK: bb2:
+// CHECK:  br bb4
+// CHECK: bb3:
+// CHECK:  return
+// CHECK: bb4
+// CHECK:  = builtin "sadd_with_overflow_Int64
+// CHECK:  br bb3
+
+sil @loop_unroll_4 : $@convention(thin) () -> () {
+bb0:
+ %0 = integer_literal $Builtin.Int64, 0
+ %1 = integer_literal $Builtin.Int64, 1
+ %2 = integer_literal $Builtin.Int64, 1
+ %3 = integer_literal $Builtin.Int1, 1
+ br bb1(%0 : $Builtin.Int64)
+
+bb1(%4 : $Builtin.Int64):
+  %5 = builtin "sadd_with_overflow_Int64"(%4 : $Builtin.Int64, %1 : $Builtin.Int64, %3 : $Builtin.Int1) : $(Builtin.Int64, Builtin.Int1)
+  %6 = tuple_extract %5 : $(Builtin.Int64, Builtin.Int1), 0
+  %7 = builtin "cmp_sgt_Int64"(%6 : $Builtin.Int64, %2 : $Builtin.Int64) : $Builtin.Int1
+  cond_br %7, bb3, bb2
+
+bb2:
+  br bb1(%6 : $Builtin.Int64)
+
+bb3:
+ %8 = tuple()
+ return %8 : $()
+}
+
 class B {}
 
 // CHECK-LABEL: sil @unroll_with_stack_allocation