[AutoDiff] Modify inlining logic to award inlining benefits to VJPs

Similar to #69029 but for VJPs.

Author: jkshtj

include/swift/SILOptimizer/Utils/PerformanceInlinerUtils.h

@@ -70,6 +70,11 @@ inline bool isOptimizableSemanticFunction(SILFunction *function) {
 /// within another semantic function, or from a "trivial" wrapper.
 bool isNestedSemanticCall(FullApplySite apply);
 
+// Strips down simple function conversion operations until a base SILValue is
+// reached.
+//
+// Returns a nullptr if `val` is not a function conversion instruction.
+SILValue stripFunctionConversions(SILValue val);
 } // end swift namespace
 
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Transforms/PerformanceInliner.cpp

@@ -213,6 +213,8 @@ class SILPerformanceInliner {
 
   bool isAutoDiffLinearMapWithControlFlow(FullApplySite AI);
 
+  bool isTupleWithAllocsOrPartialApplies(SILValue retVal);
+
   bool isProfitableToInline(
       FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
       int &NumCallerBlocks,
@@ -319,14 +321,8 @@ bool SILPerformanceInliner::isAutoDiffLinearMapWithControlFlow(
         // 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();
+      } else if (auto base = stripFunctionConversions(inst)) {
+        val = base;
         continue;
       }
       return false;
@@ -367,6 +363,29 @@ bool SILPerformanceInliner::isAutoDiffLinearMapWithControlFlow(
   return false;
 }
 
+// Checks if the given value is a tuple containing allocated objects
+// or partial applies.
+//
+// Returns true if the number of allocated objects or partial applies is
+// greater than 0, and false otherwise. 
+// 
+// Returns false if the value is not a tuple.
+bool SILPerformanceInliner::isTupleWithAllocsOrPartialApplies(SILValue val) {
+  if (auto *ti = dyn_cast<TupleInst>(val)) {
+    for (auto i : range(ti->getNumOperands())) {
+      SILValue val = ti->getOperand(i);
+
+      if (auto base = stripFunctionConversions(val))
+        val = base;
+
+      if (isa<AllocationInst>(val) || isa<PartialApplyInst>(val))
+        return true;
+    }
+  }
+
+  return false;
+}
+
 bool SILPerformanceInliner::isProfitableToInline(
     FullApplySite AI, Weight CallerWeight, ConstantTracker &callerTracker,
     int &NumCallerBlocks,
@@ -483,6 +502,9 @@ bool SILPerformanceInliner::isProfitableToInline(
         if (def && (isa<FunctionRefInst>(def) || isa<PartialApplyInst>(def)))
           BlockW.updateBenefit(Benefit, RemovedClosureBenefit);
         else if (isAutoDiffLinearMapWithControlFlow(FAI)) {
+          // TODO: Do we need to tweak inlining benefits given to pullbacks
+          // (with and without control-flow)?
+
           // 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).
@@ -587,7 +609,7 @@ bool SILPerformanceInliner::isProfitableToInline(
         // Inlining functions which return an allocated object or partial_apply
         // most likely has a benefit in the caller, because e.g. it can enable
         // de-virtualization.
-        if (isa<AllocationInst>(retVal) || isa<PartialApplyInst>(retVal)) {
+        if (isa<AllocationInst>(retVal) || isa<PartialApplyInst>(retVal) || isTupleWithAllocsOrPartialApplies(retVal)) {
           BlockW.updateBenefit(Benefit, RemovedCallBenefit + 10);
           returnsAllocation = true;
         }

lib/SILOptimizer/Utils/PerformanceInlinerUtils.cpp

@@ -117,6 +117,30 @@ static SILValue getMember(SILInstruction *inst, ProjectionPath &projStack) {
   return SILValue();
 }
 
+SILValue swift::stripFunctionConversions(SILValue val) {
+  SILValue result = nullptr;
+
+  for (;;) {
+    if (auto ti = dyn_cast<ThinToThickFunctionInst>(val)) {
+      val = ti->getOperand();
+      result = val;
+      continue;
+    } else if (auto cfi = dyn_cast<ConvertFunctionInst>(val)) {
+      val = cfi->getOperand();
+      result = val;
+      continue;
+    } else if (auto cvt = dyn_cast<ConvertEscapeToNoEscapeInst>(val)) {
+      val = cvt->getOperand();
+      result = val;
+      continue;
+    } else {
+      break;
+    }
+  }
+
+  return result;
+}
+
 SILInstruction *ConstantTracker::getDef(SILValue val,
                                         ProjectionPath &projStack) {
 
@@ -137,14 +161,8 @@ SILInstruction *ConstantTracker::getDef(SILValue val,
         // A value loaded from memory.
         val = loadedVal;
         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();
+      } else if (auto base = stripFunctionConversions(inst)) {
+        val = base;
         continue;
       }
       return inst;

test/AutoDiff/SILOptimizer/vjp_and_pullback_inlining.swift

@@ -12,6 +12,27 @@ import Glibc
 import Foundation
 #endif
 
+// ======================== VJPs ======================== //
+@differentiable(reverse)
+@_silgen_name("simple_vjp")
+func simple_vjp(x: Float) -> Float {
+    let a = x * x;
+    let b = x + x;
+    let c = x * a;
+    let d = a + b;
+    let e = b * c;
+
+    return a * b / c + d - e ;
+}
+
+@inline(never)
+@_silgen_name("caller_of_simple_vjp")
+func caller_of_simple_vjp() -> Float {
+  gradient(at: Float(4), of: simple_vjp)
+}
+
+// CHECK: decision {{{.*}}, b=30, {{.*}}} simple_vjpTJrSpSr
+// CHECK-NEXT: "simple_vjpTJrSpSr" inlined into "caller_of_simple_vjp"
 
 // ======================== Pullback w/ control-flow ======================== //