Add the Swift equivalent of the llvm.assume intrinsic

rdar://44947459

Author: aschwaighofer

include/swift/AST/Builtins.def

@@ -114,6 +114,8 @@ BUILTIN_UNARY_OPERATION(FNeg, "fneg", "n", FloatOrVector)
 // It has only an effect if the argument is a load or call.
 // TODO: consider printing a warning if it is not used on a load or call.
 BUILTIN_UNARY_OPERATION(AssumeNonNegative, "assumeNonNegative", "n", Integer)
+// It only works on i1.
+BUILTIN_UNARY_OPERATION(AssumeTrue, "assume", "", Integer)
 
 #undef BUILTIN_UNARY_OPERATION
 

lib/IRGen/GenBuiltin.cpp

@@ -393,7 +393,13 @@ if (Builtin.ID == BuiltinValueKind::id) { \
     llvm::Value *v = IGF.Builder.CreateFSub(lhs, rhs);
     return out.add(v);
   }
-  
+  if (Builtin.ID == BuiltinValueKind::AssumeTrue) {
+    llvm::Value *v = args.claimNext();
+    if (v->getType() == IGF.IGM.Int1Ty) {
+      IGF.Builder.CreateIntrinsicCall(llvm::Intrinsic::ID::assume, v);
+    }
+    return;
+  }
   if (Builtin.ID == BuiltinValueKind::AssumeNonNegative) {
     llvm::Value *v = args.claimNext();
     // Set a value range on the load instruction, which must be the argument of

lib/SIL/SILOwnershipVerifier.cpp

@@ -1157,6 +1157,7 @@ CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AssignCopyArrayFrontToBack)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AssignCopyArrayBackToFront)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AssignTakeArray)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AssumeNonNegative)
+CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AssumeTrue)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AtomicLoad)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AtomicRMW)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, MustBeLive, AtomicStore)

lib/SIL/ValueOwnershipKindClassifier.cpp

@@ -401,6 +401,7 @@ CONSTANT_OWNERSHIP_BUILTIN(Trivial, AShr)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, Add)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, And)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, AssumeNonNegative)
+CONSTANT_OWNERSHIP_BUILTIN(Trivial, AssumeTrue)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, BitCast)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, ExactSDiv)
 CONSTANT_OWNERSHIP_BUILTIN(Trivial, ExactUDiv)

stdlib/public/core/Builtin.swift

@@ -290,6 +290,11 @@ public func _onFastPath() {
   Builtin.onFastPath()
 }
 
+@_transparent
+public func _assume(_ condition: Bool) {
+  _ = Builtin.assume_Int1(condition._value)
+}
+
 //===--- Runtime shim wrappers --------------------------------------------===//
 
 /// Returns `true` iff the class indicated by `theClass` uses native

test/IRGen/assume.swift

@@ -0,0 +1,15 @@
+// RUN: %target-swift-frontend -module-name=A -O -emit-ir %s | %FileCheck %s
+
+// LLVM uses the assume intrinsic to strength reduce the division.
+
+// CHECK-LABEL: define swiftcc i64 @"$s1A10testAssumeyS2iF"(i64)
+// CHECK:   [[COND:%.*]] = icmp sgt i64 %0, -1
+// CHECK:   tail call void @llvm.assume(i1 [[COND]])
+// CHECK:   [[RES:%.*]] = lshr i64 %0, 6
+// CHECK:   ret i64 [[RES]]
+
+public func testAssume(_ i: Int) -> Int {
+  let cond = i >= 0
+  _assume(cond)
+  return i / 64
+}

test/IRGen/builtins.swift

@@ -652,6 +652,12 @@ func isUnique(_ ref: inout Builtin.BridgeObject) -> Bool {
   return Builtin.isUnique(&ref)
 }
 
+// CHECK-LABEL: define hidden{{.*}} @"$s8builtins10assumeTrueyyBi1_F"
+// CHECK: call void @llvm.assume(i1 %0)
+// CHECK: ret
+func assumeTrue(_ x: Builtin.Int1) {
+  Builtin.assume_Int1(x)
+}
 // BridgeObject nonNull
 // CHECK-LABEL: define hidden {{.*}}i1 @"$s8builtins15isUnique_nativeyBi1_BbzF"(%swift.bridge** nocapture dereferenceable({{.*}})) {{.*}} {
 // CHECK-NEXT: entry:

test/SILGen/builtins.swift

@@ -812,4 +812,10 @@ func once(control: Builtin.RawPointer) {
 func valueToBridgeObject(_ x: UInt) -> Builtin.BridgeObject {
   return Builtin.valueToBridgeObject(x)
 }
- 
+
+// CHECK-LABEL: sil hidden @$s8builtins10assumeTrueyyBi1_F
+// CHECK: builtin "assume_Int1"({{.*}} : $Builtin.Int1)
+// CHECK: return
+func assumeTrue(_ x: Builtin.Int1) {
+  Builtin.assume_Int1(x)
+}