SIL: branch weights for try_apply's

kavon · kavon · commit 0cce602bce84 · 2024-08-08T21:21:52.000-04:00
diff --git a/include/swift/SIL/SILBuilder.h b/include/swift/SIL/SILBuilder.h
@@ -555,10 +555,13 @@ class SILBuilder {
       ArrayRef<SILValue> args, SILBasicBlock *normalBB, SILBasicBlock *errorBB,
       ApplyOptions options = ApplyOptions(),
       const GenericSpecializationInformation *specializationInfo = nullptr,
-      std::optional<ApplyIsolationCrossing> isolationCrossing = std::nullopt) {
+      std::optional<ApplyIsolationCrossing> isolationCrossing = std::nullopt,
+      ProfileCounter normalCount = ProfileCounter(),
+      ProfileCounter errorCount = ProfileCounter()) {
     return insertTerminator(TryApplyInst::create(
         getSILDebugLocation(loc), callee, subs, args, normalBB, errorBB,
-        options, *F, specializationInfo, isolationCrossing));
+        options, *F, specializationInfo, isolationCrossing,
+        normalCount, errorCount));
   }
 
   PartialApplyInst *createPartialApply(
diff --git a/include/swift/SIL/SILInstruction.h b/include/swift/SIL/SILInstruction.h
@@ -10968,7 +10968,8 @@ class TryApplyInstBase : public TermInst {
 
 protected:
   TryApplyInstBase(SILInstructionKind valueKind, SILDebugLocation Loc,
-                   SILBasicBlock *normalBB, SILBasicBlock *errorBB);
+                   SILBasicBlock *normalBB, SILBasicBlock *errorBB,
+                   ProfileCounter normalCount, ProfileCounter errorCount);
 
 public:
   SuccessorListTy getSuccessors() {
@@ -10988,6 +10989,11 @@ class TryApplyInstBase : public TermInst {
   const SILBasicBlock *getNormalBB() const { return DestBBs[NormalIdx]; }
   SILBasicBlock *getErrorBB() { return DestBBs[ErrorIdx]; }
   const SILBasicBlock *getErrorBB() const { return DestBBs[ErrorIdx]; }
+
+  /// The number of times the Normal branch was executed
+  ProfileCounter getNormalBBCount() const { return DestBBs[NormalIdx].getCount(); }
+  /// The number of times the Error branch was executed
+  ProfileCounter getErrorBBCount() const { return DestBBs[ErrorIdx].getCount(); }
 };
 
 /// TryApplyInst - Represents the full application of a function that
@@ -11005,15 +11011,19 @@ class TryApplyInst final
                SILBasicBlock *normalBB, SILBasicBlock *errorBB,
                ApplyOptions options,
                const GenericSpecializationInformation *specializationInfo,
-               std::optional<ApplyIsolationCrossing> isolationCrossing);
+               std::optional<ApplyIsolationCrossing> isolationCrossing,
+               ProfileCounter normalCount,
+               ProfileCounter errorCount);
 
   static TryApplyInst *
   create(SILDebugLocation debugLoc, SILValue callee,
          SubstitutionMap substitutions, ArrayRef<SILValue> args,
          SILBasicBlock *normalBB, SILBasicBlock *errorBB, ApplyOptions options,
          SILFunction &parentFunction,
          const GenericSpecializationInformation *specializationInfo,
-         std::optional<ApplyIsolationCrossing> isolationCrossing);
+         std::optional<ApplyIsolationCrossing> isolationCrossing,
+         ProfileCounter normalCount,
+         ProfileCounter errorCount);
 };
 
 /// DifferentiableFunctionInst - creates a `@differentiable` function-typed
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -3884,9 +3884,18 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
     // FIXME: Remove this when the following radar is fixed: rdar://116636601
     Builder.CreatePtrToInt(errorValue, IGM.IntPtrTy);
 
+    // Emit profile metadata if available.
+    llvm::MDNode *Weights = nullptr;
+    auto NormalBBCount = tryApplyInst->getNormalBBCount();
+    auto ErrorBBCount = tryApplyInst->getErrorBBCount();
+    if (NormalBBCount || ErrorBBCount)
+      Weights = IGM.createProfileWeights(ErrorBBCount ? ErrorBBCount.getValue() : 0,
+                                         NormalBBCount ? NormalBBCount.getValue() : 0);
+
     Builder.CreateCondBr(hasError,
                          typedErrorLoadBB ? typedErrorLoadBB : errorDest.bb,
-                         normalDest.bb);
+                         normalDest.bb,
+                         Weights);
 
     // Set up the PHI nodes on the normal edge.
     unsigned firstIndex = 0;
diff --git a/lib/SIL/IR/SILInstructions.cpp b/lib/SIL/IR/SILInstructions.cpp
@@ -783,19 +783,24 @@ PartialApplyInst *PartialApplyInst::create(
 TryApplyInstBase::TryApplyInstBase(SILInstructionKind kind,
                                    SILDebugLocation loc,
                                    SILBasicBlock *normalBB,
-                                   SILBasicBlock *errorBB)
-    : TermInst(kind, loc), DestBBs{{{this, normalBB}, {this, errorBB}}} {}
+                                   SILBasicBlock *errorBB,
+                                   ProfileCounter normalCount,
+                                   ProfileCounter errorCount)
+    : TermInst(kind, loc), DestBBs{{{this, normalBB, normalCount},
+                                    {this, errorBB, errorCount}}} {}
 
 TryApplyInst::TryApplyInst(
     SILDebugLocation loc, SILValue callee, SILType substCalleeTy,
     SubstitutionMap subs, ArrayRef<SILValue> args,
     ArrayRef<SILValue> typeDependentOperands, SILBasicBlock *normalBB,
     SILBasicBlock *errorBB, ApplyOptions options,
     const GenericSpecializationInformation *specializationInfo,
-    std::optional<ApplyIsolationCrossing> isolationCrossing)
+    std::optional<ApplyIsolationCrossing> isolationCrossing,
+    ProfileCounter normalCount,
+    ProfileCounter errorCount)
     : InstructionBase(isolationCrossing, loc, callee, substCalleeTy, subs, args,
                       typeDependentOperands, specializationInfo, normalBB,
-                      errorBB) {
+                      errorBB, normalCount, errorCount) {
   setApplyOptions(options);
 }
 
@@ -805,19 +810,33 @@ TryApplyInst::create(SILDebugLocation loc, SILValue callee,
                      SILBasicBlock *normalBB, SILBasicBlock *errorBB,
                      ApplyOptions options, SILFunction &parentFunction,
                      const GenericSpecializationInformation *specializationInfo,
-                     std::optional<ApplyIsolationCrossing> isolationCrossing) {
+                     std::optional<ApplyIsolationCrossing> isolationCrossing,
+                     ProfileCounter normalCount,
+                     ProfileCounter errorCount) {
   SILType substCalleeTy = callee->getType().substGenericArgs(
       parentFunction.getModule(), subs,
       parentFunction.getTypeExpansionContext());
 
+  if (parentFunction.getModule().getOptions().EnableThrowsPrediction &&
+      !normalCount && !errorCount) {
+    // Predict that the error branch is not taken.
+    //
+    // We cannot use the Expect builtin within SIL because try_apply abstracts
+    // over the raw conditional test to see if an error was returned.
+    // So, we synthesize profiling branch weights instead.
+    normalCount = 1999;
+    errorCount = 0;
+  }
+
   SmallVector<SILValue, 32> typeDependentOperands;
   collectTypeDependentOperands(typeDependentOperands, parentFunction,
                                substCalleeTy.getASTType(), subs);
   void *buffer = allocateTrailingInst<TryApplyInst, Operand>(
       parentFunction, getNumAllOperands(args, typeDependentOperands));
   return ::new (buffer) TryApplyInst(
       loc, callee, substCalleeTy, subs, args, typeDependentOperands, normalBB,
-      errorBB, options, specializationInfo, isolationCrossing);
+      errorBB, options, specializationInfo, isolationCrossing,
+      normalCount, errorCount);
 }
 
 SILType DifferentiableFunctionInst::getDifferentiableFunctionType(
diff --git a/lib/SIL/IR/SILPrinter.cpp b/lib/SIL/IR/SILPrinter.cpp
@@ -1573,6 +1573,10 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
     visitApplyInstBase(AI);
     *this << ", normal " << Ctx.getID(AI->getNormalBB());
     *this << ", error " << Ctx.getID(AI->getErrorBB());
+    if (AI->getNormalBBCount())
+      *this << " !normal_count(" << AI->getNormalBBCount().getValue() << ")";
+    if (AI->getErrorBBCount())
+      *this << " !error_count(" << AI->getErrorBBCount().getValue() << ")";
   }
 
   void visitPartialApplyInst(PartialApplyInst *CI) {
diff --git a/test/SILOptimizer/throws_prediction.swift b/test/SILOptimizer/throws_prediction.swift
@@ -0,0 +1,122 @@
+// RUN: %target-swift-frontend %s \
+// RUN:   -disable-availability-checking \
+// RUN:   -enable-throws-prediction \
+// RUN:   -sil-verify-all -module-name=test -emit-sil \
+// RUN:       | %FileCheck --check-prefix CHECK-SIL %s
+
+// RUN: %target-swift-frontend %s \
+// RUN:   -disable-availability-checking \
+// RUN:   -enable-throws-prediction \
+// RUN:   -sil-verify-all -module-name=test -emit-irgen \
+// RUN:       | %FileCheck --check-prefix CHECK-IR %s
+
+// RUN: %target-swift-frontend %s \
+// RUN:   -disable-availability-checking \
+// RUN:   -disable-throws-prediction \
+// RUN:   -sil-verify-all -module-name=test -emit-sil \
+// RUN:       | %FileCheck --check-prefix CHECK-DISABLED %s
+
+// CHECK-DISABLED-NOT: normal_count
+
+enum MyError: Error { case err }
+
+func throwy1() throws {}
+func throwy2() throws(MyError) { }
+func throwy3() async throws -> Int { 0 }
+func throwy4() async throws(MyError) -> Int { 1 }
+
+// CHECK-SIL-LABEL: sil hidden @$s4test0A13TryPredictionyySbF
+// CHECK-SIL: try_apply {{.*}} @error any Error{{.*}} !normal_count(1999) !error_count(0)
+// CHECK-SIL: try_apply {{.*}} @error MyError{{.*}} !normal_count(1999) !error_count(0)
+
+// CHECK-IR-LABEL: define hidden swiftcc void @"$s4test0A13TryPredictionyySbF"
+// CHECK-IR:   call swiftcc void @"$s4test7throwy1yyKF"
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+
+// CHECK-IR:   call swiftcc void @"$s4test7throwy2yyAA7MyErrorOYKF"
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+func testTryPrediction(_ b: Bool) {
+  do {
+    try throwy1()
+    try throwy2()
+  } catch {
+    print("hi")
+  }
+}
+
+// CHECK-SIL-LABEL: sil hidden @$s4test0A21AsyncThrowsPredictionSiyYaF
+// CHECK-SIL: function_ref @$s4test7throwy3SiyYaKF
+// CHECK-SIL: try_apply {{.*}} @error any Error{{.*}} !normal_count(1999) !error_count(0)
+
+// CHECK-IR-LABEL: define hidden swifttailcc void @"$s4test0A21AsyncThrowsPredictionSiyYaF"
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+func testAsyncThrowsPrediction() async -> Int {
+    if let x = try? await throwy3() {
+        return x
+    }
+    return 1337
+}
+
+// CHECK-SIL-LABEL: sil hidden @$s4test0A28Async_TYPED_ThrowsPredictionSiyYaAA7MyErrorOYKF
+// CHECK-SIL: try_apply {{.*}} @error MyError{{.*}} !normal_count(1999) !error_count(0)
+// CHECK-SIL: try_apply {{.*}} @error MyError{{.*}} !normal_count(1999) !error_count(0)
+// CHECK-SIL: try_apply {{.*}} @error MyError{{.*}} !normal_count(1999) !error_count(0)
+
+// CHECK-IR-LABEL: define hidden swifttailcc void @"$s4test0A28Async_TYPED_ThrowsPredictionSiyYaAA7MyErrorOYKF"
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+//
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+//
+// CHECK-IR:   [[ERROR_PTR:%.*]] = load ptr, ptr %swifterror
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr [[ERROR_PTR]], null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+func testAsync_TYPED_ThrowsPrediction() async throws(MyError) -> Int {
+    let x = try await throwy4()
+    let y = try await throwy4()
+    let z = try await throwy4()
+    return x + y + z
+}
+
+
+func getRandom(_ b: Bool) throws -> Int {
+    if b {
+        return Int.random(in: 0..<1024)
+    } else {
+        throw MyError.err
+    }
+}
+
+// CHECK-SIL-LABEL: sil hidden @$s4test20sequenceOfNormalTrysySiSb_S2btKF
+// CHECK-SIL: try_apply {{.*}} @error any Error{{.*}} !normal_count(1999) !error_count(0)
+// CHECK-SIL: try_apply {{.*}} @error any Error{{.*}} !normal_count(1999) !error_count(0)
+// CHECK-SIL: try_apply {{.*}} @error any Error{{.*}} !normal_count(1999) !error_count(0)
+
+// CHECK-IR-LABEL: define hidden swiftcc i64 @"$s4test20sequenceOfNormalTrysySiSb_S2btKF"
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr {{%.*}}, null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+//
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr {{%.*}}, null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+//
+// CHECK-IR:   [[HAVE_ERROR:%.*]] = icmp ne ptr {{%.*}}, null
+// CHECK-IR:   br i1 [[HAVE_ERROR]], {{.*}} !prof [[PREFER_FALSE:![0-9]+]]
+func sequenceOfNormalTrys(_ b1: Bool,
+                          _ b2: Bool,
+                          _ b3: Bool) throws -> Int {
+    let x = try getRandom(b1)
+    let y = try getRandom(b2)
+    let z = try getRandom(b3)
+    return x + y + z
+}
+
+// CHECK-IR: [[PREFER_FALSE]] = !{!"branch_weights", i32 1, i32 2000}
