[Autodiff] Modify inliner logic to award inlining benefits to linear maps w/ control-flow

For linear maps containing control-flow, closures (representing the pullbacks
of intermediate values) may be passed as arguments, however, they may be
hidden behind a branch-tracing enum (tracing execution flow of the original
function).

Such linear maps did not use to get inlining benefits as the compiler
could not see that the intermediate pullback closures were actually part
of the input.

This change modifies the inliner logic to correctly award inlining
benefits to linear maps containing control-flow, by checking if a
"callee" in the linear map actually traces back to an input closure that
was received as part of a branch-tracing enum input argument.

Fixes #68945

Author: jkshtj

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -208,6 +208,8 @@ class SILPerformanceInliner {
           llvm::detail::DenseMapPair<swift::SILBasicBlock *, uint64_t>, true>
           &bbIt);
 
+  bool isAutoDiffLinearMapWithControlFlow(FullApplySite AI);
+
   bool isProfitableToInline(
       FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
       int &NumCallerBlocks,
@@ -299,6 +301,69 @@ bool SILPerformanceInliner::profileBasedDecision(
   return true;
 }
 
+// Checks if `FAI` can be traced back to a specifically named,
+// input enum function argument. If so, the callsite
+// containing function is a linear map in Swift Autodiff.
+bool SILPerformanceInliner::isAutoDiffLinearMapWithControlFlow(
+    FullApplySite FAI) {
+  static const std::string LinearMapBranchTracingEnumPrefix = "_AD__";
+
+  auto val = FAI.getCallee();
+
+  for (;;) {
+    if (auto *inst = dyn_cast<SingleValueInstruction>(val)) {
+      if (auto pi = Projection::isObjectProjection(val)) {
+        // Extract a member from a struct/tuple/enum.
+        val = pi->getOperand(0);
+        continue;
+      } else if (auto ti = dyn_cast<ThinToThickFunctionInst>(inst)) {
+        val = ti->getOperand();
+        continue;
+      } else if (auto cfi = dyn_cast<ConvertFunctionInst>(inst)) {
+        val = cfi->getOperand();
+        continue;
+      } else if (auto cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(inst)) {
+        val = cvt->getOperand();
+        continue;
+      }
+      return false;
+    } else if (auto *phiArg = dyn_cast<SILPhiArgument>(val)) {
+      if (auto *predBB = phiArg->getParent()->getSinglePredecessorBlock()) {
+        // The terminator of this predecessor block must either be a
+        // (conditional) branch instruction or a switch_enum.
+        if (auto *bi = dyn_cast<BranchInst>(predBB->getTerminator())) {
+          val = bi->getArg(phiArg->getIndex());
+          continue;
+        } else if (auto *cbi =
+                       dyn_cast<CondBranchInst>(predBB->getTerminator())) {
+          val = cbi->getArgForDestBB(phiArg->getParent(), phiArg->getIndex());
+          continue;
+        } else if (auto *sei =
+                       dyn_cast<SwitchEnumInst>(predBB->getTerminator())) {
+          val = sei->getOperand();
+          continue;
+        }
+        return false;
+      }
+    }
+    break;
+  }
+
+  // If `val` now points to a function argument then we have successfully traced
+  // the callee back to a function argument.
+  //
+  // We now need to check if this argument is an enum and named like an autodiff
+  // branch tracing enum.
+  if (auto *arg = dyn_cast<SILFunctionArgument>(val)) {
+    if (auto *enumDecl = arg->getType().getEnumOrBoundGenericEnum()) {
+      return enumDecl->getName().str().startswith(
+          LinearMapBranchTracingEnumPrefix);
+    }
+  }
+
+  return false;
+}
+
 bool SILPerformanceInliner::isProfitableToInline(
     FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
     int &NumCallerBlocks,
@@ -413,6 +478,19 @@ bool SILPerformanceInliner::isProfitableToInline(
         SILInstruction *def = constTracker.getDefInCaller(FAI.getCallee());
         if (def && (isa<FunctionRefInst>(def) || isa<PartialApplyInst>(def)))
           BlockW.updateBenefit(Benefit, RemovedClosureBenefit);
+        else if (isAutoDiffLinearMapWithControlFlow(FAI)) {
+          // For linear maps in Swift Autodiff, callees may be passed as an
+          // argument, however, they may be hidden behind a branch-tracing
+          // enum (tracing execution flow of the original function).
+          //
+          // If we can establish that we are inside of a Swift Autodiff linear
+          // map and that the branch tracing input enum is wrapping pullback
+          // closures, then we can update this function's benefit with
+          // `RemovedClosureBenefit` because inlining will (probably) eliminate
+          // the closure.
+          BlockW.updateBenefit(Benefit, RemovedClosureBenefit);
+        }
+
         // Check if inlining the callee would allow for further
         // optimizations like devirtualization or generic specialization. 
         if (!def)

test/AutoDiff/SILOptimizer/vjp_and_pullback_inlining.swift

@@ -0,0 +1,79 @@
+// VJP and pullback inlining tests.
+
+// RUN: %target-swift-frontend -emit-sil -O -verify -Xllvm -debug-only=sil-inliner %s 2>&1 | %FileCheck %s
+
+// REQUIRES: asserts
+// REQUIRES: swift_in_compiler
+
+import _Differentiation
+#if os(Linux)
+import Glibc
+#else
+import Foundation
+#endif
+
+
+// ======================== Pullback w/ control-flow ======================== //
+
+@differentiable(reverse)
+@_silgen_name("pb_with_control_flow")
+func pb_with_control_flow(_ x: Float) -> Float {
+  if (x > 0) {
+    return sin(x) * cos(x)
+  } else {
+    return sin(x) + cos(x)
+  }
+}
+
+@inline(never)
+@_silgen_name("caller_of_pb_with_control_flow")
+func caller_of_pb_with_control_flow() -> Float {
+    gradient(at: Float(1), of: pb_with_control_flow)
+}
+
+// CHECK: decision {{{.*}}, b=70, {{.*}}} pb_with_control_flowTJpSpSr
+// CHECK-NEXT: "pb_with_control_flowTJpSpSr" inlined into "caller_of_pb_with_control_flow"
+
+
+@differentiable(reverse)
+func double(x: Float) -> Float {
+    return x + x
+}
+
+@differentiable(reverse)
+func square(x: Float) -> Float {
+    return x * x
+}
+
+@differentiable(reverse)
+@_silgen_name("more_complex_pb_with_control_flow")
+func more_complex_pb_with_control_flow(x: Float) -> Float {
+    if (x > 0) {
+        if ((x+1) < 5) {
+            if (x*2 > 4) {
+                let y = square(x: x)
+                if (y >= x) {
+                    let d = double(x: x)
+                    return x - (d*y)
+                } else {
+                    let e = square(x: y)
+                    return x + (e*y)
+                }
+            }
+        }
+    } else {
+        let y = double(x: x)
+        return x * y
+    }
+
+    return x*3
+}
+
+@inline(never)
+@_silgen_name("caller_of_more_complex_pb_with_control_flow")
+func caller_of_more_complex_pb_with_control_flow() -> Float {
+    gradient(at: Float(1), of: more_complex_pb_with_control_flow)
+}
+
+// CHECK: decision {{{.*}}, b=70, {{.*}}} more_complex_pb_with_control_flowTJpSpSr
+// CHECK-NEXT: "more_complex_pb_with_control_flowTJpSpSr" inlined into "caller_of_more_complex_pb_with_control_flow"