IRGen: better code generation for the _assumeNonNegative intrinsic

So far we only supported this for `load` and `call` arguments. Now we use `llvm.assume` for other kind of argument values.
This helps optimizing bounds checking for `Span`.

Author: eeckstein

lib/IRGen/GenBuiltin.cpp

@@ -600,8 +600,8 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
   }
   if (Builtin.ID == BuiltinValueKind::AssumeNonNegative) {
     llvm::Value *v = args.claimNext();
-    // Set a value range on the load instruction, which must be the argument of
-    // the builtin.
+    // If the argument is a `load` or `call` we can use a range metadata to
+    // specify the >= 0 constraint.
     if (isa<llvm::LoadInst>(v) || isa<llvm::CallInst>(v)) {
       // The load must be post-dominated by the builtin. Otherwise we would get
       // a wrong assumption in the else-branch in this example:
@@ -626,6 +626,10 @@ void irgen::emitBuiltinCall(IRGenFunction &IGF, const BuiltinInfo &Builtin,
         llvm::MDNode *range = llvm::MDNode::get(ctx, rangeElems);
         I->setMetadata(llvm::LLVMContext::MD_range, range);
       }
+    } else {
+      // Otherwise, we specify the constraint with an `llvm.assume` intrinsic.
+      auto *cmp = IGF.Builder.CreateICmpSGE(v, llvm::ConstantInt::get(v->getType(), 0));
+      IGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::assume, cmp);
     }
     // Don't generate any code for the builtin.
     return out.add(v);

test/IRGen/builtins.swift

@@ -646,6 +646,13 @@ func assumeNonNegative_test2(_ x: Builtin.Word) -> Builtin.Word {
   return Builtin.assumeNonNegative_Word(return_word(x))
 }
 
+// CHECK-LABEL: define hidden {{.*}}@"$s8builtins23assumeNonNegative_test3yBwBwF"
+func assumeNonNegative_test3(_ x: Builtin.Word) -> Builtin.Word {
+  // CHECK: [[C:%.*]] = icmp sge i{{[0-9]+}} %0, 0
+  // CHECK: call void @llvm.assume(i1 [[C]])
+  return Builtin.assumeNonNegative_Word(x)
+}
+
 struct Empty {}
 struct Pair { var i: Int, b: Bool }
 

test/SILOptimizer/span_bounds_check_tests.swift

@@ -28,7 +28,7 @@ public func span_sum_iterate_to_count_wo_trap(_ v: Span<Int>) -> Int {
   for i in 0..<v.count {
     sum &+= v[i]
   }
-  return sum
+  return sum &+ 1 // adding 1 to avoid LLVM function merging
 }
 
 // Bounds check should be eliminated
@@ -175,15 +175,10 @@ public func span_sum_iterate_to_deducible_count2_with_trap(_ v: Span<Int>, _ n:
 }
 
 // Bounds check should be eliminated
-// SIL optimizer hoists bounds checks from the loop
+// TODO: there should not be any cond_fail in this function
 
 // CHECK-SIL-LABEL: sil @$s23span_bounds_check_tests0A29_iterate_over_indices_wo_trapySis4SpanVySiGF :
-// CHECK-SIL: bb1
-// CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: bb3
 // CHECK-SIL-NOT: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: cond_br
 // CHECK-SIL-LABEL: } // end sil function '$s23span_bounds_check_tests0A29_iterate_over_indices_wo_trapySis4SpanVySiGF'
 
 // CHECK-IR: define {{.*}} @"$s23span_bounds_check_tests0A29_iterate_over_indices_wo_trapySis4SpanVySiGF"
@@ -200,12 +195,7 @@ public func span_iterate_over_indices_wo_trap(_ v: Span<Int>) -> Int {
 // SIL optimizer hoists bounds checks from the loop
 
 // CHECK-SIL-LABEL: sil @$s23span_bounds_check_tests0A31_iterate_over_indices_with_trapySis4SpanVySiGF :
-// CHECK-SIL: bb1
-// CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: bb3
 // CHECK-SIL-NOT: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: cond_br
 // CHECK-SIL-LABEL: } // end sil function '$s23span_bounds_check_tests0A31_iterate_over_indices_with_trapySis4SpanVySiGF'
 public func span_iterate_over_indices_with_trap(_ v: Span<Int>) -> Int {
   var sum = 0
@@ -254,12 +244,11 @@ public func span_search<T : Equatable>(_ v: Span<T>, _ elem: T) -> Int? {
   return nil
 }
 
-// Bounds check should be eliminated
+// TODO: Bounds check should be eliminated
 
 // CHECK-SIL-LABEL: sil @$s23span_bounds_check_tests0A11_search_splySiSgs4SpanVySiG_SitF :
 // CHECK-SIL: bb3:
 // CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
-// CHECK-SIL: cond_fail {{.*}}, "Index out of bounds"
 // CHECK-SIL: cond_br
 // CHECK-SIL-LABEL: } // end sil function '$s23span_bounds_check_tests0A11_search_splySiSgs4SpanVySiG_SitF'
 public func span_search_spl(_ v: Span<Int>, _ elem: Int) -> Int? {