Merge pull request swiftlang#83172 from eeckstein/span-bounds-checking

Improve Span bounds check elimination by adding `_assumeNonNegative` to `count`

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);

stdlib/public/core/Span/MutableSpan.swift

@@ -237,7 +237,7 @@ extension MutableSpan where Element: ~Copyable {
 extension MutableSpan where Element: ~Copyable {
 
   @_alwaysEmitIntoClient
-  public var count: Int { _count }
+  public var count: Int { _assumeNonNegative(_count) }
 
   @_alwaysEmitIntoClient
   public var isEmpty: Bool { _count == 0 }
@@ -246,7 +246,7 @@ extension MutableSpan where Element: ~Copyable {
 
   @_alwaysEmitIntoClient
   public var indices: Range<Index> {
-    unsafe Range(_uncheckedBounds: (0, _count))
+    unsafe Range(_uncheckedBounds: (0, count))
   }
 }
 
@@ -290,18 +290,12 @@ extension MutableSpan where Element: ~Copyable {
   @_alwaysEmitIntoClient
   public subscript(_ position: Index) -> Element {
     unsafeAddress {
-      _precondition(
-        UInt(bitPattern: position) < UInt(bitPattern: _count),
-        "Index out of bounds"
-      )
+      _precondition(indices.contains(position), "index out of bounds")
       return unsafe UnsafePointer(_unsafeAddressOfElement(unchecked: position))
     }
     @lifetime(self: copy self)
     unsafeMutableAddress {
-      _precondition(
-        UInt(bitPattern: position) < UInt(bitPattern: _count),
-        "Index out of bounds"
-      )
+      _precondition(indices.contains(position), "index out of bounds")
        return unsafe _unsafeAddressOfElement(unchecked: position)
     }
   }

stdlib/public/core/Span/Span.swift

@@ -391,7 +391,7 @@ extension Span where Element: ~Copyable {
   /// - Complexity: O(1)
   @_alwaysEmitIntoClient
   @_semantics("fixed_storage.get_count")
-  public var count: Int { _count }
+  public var count: Int { _assumeNonNegative(_count) }
 
   /// A Boolean value indicating whether the span is empty.
   ///
@@ -408,7 +408,7 @@ extension Span where Element: ~Copyable {
   /// - Complexity: O(1)
   @_alwaysEmitIntoClient
   public var indices: Range<Index> {
-    unsafe Range(_uncheckedBounds: (0, _count))
+    unsafe Range(_uncheckedBounds: (0, count))
   }
 }
 
@@ -419,10 +419,7 @@ extension Span where Element: ~Copyable {
   @inline(__always)
   @_alwaysEmitIntoClient
   internal func _checkIndex(_ position: Index) {
-    _precondition(
-      UInt(bitPattern: position) < UInt(bitPattern: _count),
-      "Index out of bounds"
-    )
+    _precondition(indices.contains(position), "Index out of bounds")
   }
 
   /// Accesses the element at the specified position in the `Span`.

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? {