Merge branch 'tensorflow' into marcrasi-use-diff-witness

Author: dan-zheng

include/swift/AST/Attr.def

@@ -514,20 +514,21 @@ DECL_ATTR(differentiable, Differentiable,
   91)
 DECL_ATTR(differentiating, Differentiating,
   OnFunc | LongAttribute | AllowMultipleAttributes |
-  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove |
+  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove |
   NotSerialized, 92)
 SIMPLE_DECL_ATTR(compilerEvaluable, CompilerEvaluable,
   OnAccessor | OnFunc | OnConstructor | OnSubscript |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove |
   NotSerialized, 93)
 SIMPLE_DECL_ATTR(noDerivative, NoDerivative,
   OnVar |
-  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
+  ABIBreakingToAdd | ABIBreakingToRemove | APIBreakingToAdd | APIBreakingToRemove,
   94)
 DECL_ATTR(transposing, Transposing,
   OnFunc | LongAttribute | AllowMultipleAttributes |
-  ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove |
+  ABIStableToAdd | ABIBreakingToRemove | APIStableToAdd | APIBreakingToRemove |
   NotSerialized, 96)
+// SWIFT_ENABLE_TENSORFLOW END
 
 SIMPLE_DECL_ATTR(IBSegueAction, IBSegueAction,
   OnFunc |

include/swift/AST/DiagnosticsSIL.def

@@ -524,6 +524,8 @@ NOTE(autodiff_when_differentiating_function_definition,none,
      "when differentiating this function definition", ())
 NOTE(autodiff_cannot_differentiate_through_inout_arguments,none,
      "cannot differentiate through 'inout' arguments", ())
+NOTE(autodiff_cannot_differentiate_through_multiple_results,none,
+     "cannot differentiate through multiple results", ())
 NOTE(autodiff_class_member_not_supported,none,
      "differentiating class members is not yet supported", ())
 // TODO(TF-642): Remove when `partial_apply` works with `@differentiable`

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1513,20 +1513,6 @@ class DifferentiableActivityInfo {
                        const SILAutoDiffIndices &indices) const;
 };
 
-/// Given a parameter argument (not indirect result) and some differentiation
-/// indices, figure out whether the parent function is being differentiated with
-/// respect to this parameter, according to the indices.
-static bool isDifferentiationParameter(SILArgument *argument,
-                                       IndexSubset *indices) {
-  if (!argument) return false;
-  auto *function = argument->getFunction();
-  auto paramArgs = function->getArgumentsWithoutIndirectResults();
-  for (unsigned i : indices->getIndices())
-    if (paramArgs[i] == argument)
-      return true;
-  return false;
-}
-
 /// For an `apply` instruction with active results, compute:
 /// - The results of the `apply` instruction, in type order.
 /// - The set of minimal parameter and result indices for differentiating the
@@ -1543,9 +1529,7 @@ static void collectMinimalIndicesForFunctionCall(
   // Record all parameter indices in type order.
   unsigned currentParamIdx = 0;
   for (auto applyArg : ai->getArgumentsWithoutIndirectResults()) {
-    if (activityInfo.isVaried(applyArg, parentIndices.parameters) ||
-        isDifferentiationParameter(dyn_cast<SILArgument>(applyArg),
-                                   parentIndices.parameters))
+    if (activityInfo.isActive(applyArg, parentIndices))
       paramIndices.push_back(currentParamIdx);
     ++currentParamIdx;
   }
@@ -1566,13 +1550,12 @@ static void collectMinimalIndicesForFunctionCall(
     if (res.isFormalDirect()) {
       results.push_back(directResults[dirResIdx]);
       if (auto dirRes = directResults[dirResIdx])
-        if (dirRes && activityInfo.isUseful(dirRes, parentIndices.source))
+        if (dirRes && activityInfo.isActive(dirRes, parentIndices))
           resultIndices.push_back(idx);
       ++dirResIdx;
     } else {
       results.push_back(indirectResults[indResIdx]);
-      if (activityInfo.isUseful(indirectResults[indResIdx],
-                                parentIndices.source))
+      if (activityInfo.isActive(indirectResults[indResIdx], parentIndices))
         resultIndices.push_back(idx);
       ++indResIdx;
     }
@@ -3813,10 +3796,12 @@ class VJPEmitter final
                    activeResultIndices.begin(), activeResultIndices.end(),
                    [&s](unsigned i) { s << i; }, [&s] { s << ", "; });
                s << "}\n";);
-    // FIXME: We don't support multiple active results yet.
+    // Diagnose multiple active results.
+    // TODO(TF-983): Support multiple active results.
     if (activeResultIndices.size() > 1) {
       context.emitNondifferentiabilityError(
-          ai, invoker, diag::autodiff_expression_not_differentiable_note);
+          ai, invoker,
+          diag::autodiff_cannot_differentiate_through_multiple_results);
       errorOccurred = true;
       return;
     }
@@ -5492,13 +5477,16 @@ class JVPEmitter final
                    activeResultIndices.begin(), activeResultIndices.end(),
                    [&s](unsigned i) { s << i; }, [&s] { s << ", "; });
                s << "}\n";);
-    // FIXME: We don't support multiple active results yet.
+    // Diagnose multiple active results.
+    // TODO(TF-983): Support multiple active results.
     if (activeResultIndices.size() > 1) {
       context.emitNondifferentiabilityError(
-          ai, invoker, diag::autodiff_expression_not_differentiable_note);
+          ai, invoker,
+          diag::autodiff_cannot_differentiate_through_multiple_results);
       errorOccurred = true;
       return;
     }
+
     // Form expected indices, assuming there's only one result.
     SILAutoDiffIndices indices(
         activeResultIndices.front(),
@@ -6257,6 +6245,8 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
     // Adjoint values of dominated active values are passed as pullback block
     // arguments.
     DominanceOrder domOrder(original.getEntryBlock(), domInfo);
+    // Keep track of visited values.
+    SmallPtrSet<SILValue, 8> visited;
     while (auto *bb = domOrder.getNext()) {
       auto &bbActiveValues = activeValues[bb];
       // If the current block has an immediate dominator, append the immediate
@@ -6266,13 +6256,12 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
         bbActiveValues.append(domBBActiveValues.begin(),
                               domBBActiveValues.end());
       }
-      SmallPtrSet<SILValue, 8> visited(bbActiveValues.begin(),
-                                       bbActiveValues.end());
-      // Register a value as active if it has not yet been visited.
       bool diagnosedActiveEnumValue = false;
-      auto addActiveValue = [&](SILValue v) {
+      // Mark the activity of a value if it has not yet been visited.
+      auto markValueActivity = [&](SILValue v) {
         if (visited.count(v))
           return;
+        visited.insert(v);
         // Diagnose active enum values. Differentiation of enum values requires
         // special adjoint value handling and is not yet supported. Diagnose
         // only the first active enum value to prevent too many diagnostics.
@@ -6288,17 +6277,19 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
         // become projections into their adjoint base buffer.
         if (Projection::isAddressProjection(v))
           return;
-        visited.insert(v);
         bbActiveValues.push_back(v);
       };
-      // Register bb arguments and all instruction operands/results.
+      // Visit bb arguments and all instruction operands/results.
+      for (auto *arg : bb->getArguments())
+        if (getActivityInfo().isActive(arg, getIndices()))
+          markValueActivity(arg);
       for (auto &inst : *bb) {
         for (auto op : inst.getOperandValues())
           if (getActivityInfo().isActive(op, getIndices()))
-            addActiveValue(op);
+            markValueActivity(op);
         for (auto result : inst.getResults())
           if (getActivityInfo().isActive(result, getIndices()))
-            addActiveValue(result);
+            markValueActivity(result);
       }
       domOrder.pushChildren(bb);
       if (errorOccurred)

lib/Serialization/DeserializeSIL.cpp

@@ -799,7 +799,9 @@ SILDeserializer::readSILFunctionChecked(DeclID FID, SILFunction *existingFn,
   // SIL_VTABLE or SIL_GLOBALVAR or SIL_WITNESS_TABLE record also means the end
   // of this SILFunction.
   while (kind != SIL_FUNCTION && kind != SIL_VTABLE && kind != SIL_GLOBALVAR &&
-         kind != SIL_WITNESS_TABLE) {
+         // SWIFT_ENABLE_TENSORFLOW
+         kind != SIL_WITNESS_TABLE && kind != SIL_DIFFERENTIABILITY_WITNESS) {
+         // SWIFT_ENABLE_TENSORFLOW END
     if (kind == SIL_BASIC_BLOCK)
       // Handle a SILBasicBlock record.
       CurrentBB = readSILBasicBlock(fn, CurrentBB, scratch);

lib/Serialization/SerializeSIL.cpp

@@ -201,6 +201,12 @@ namespace {
     /// Global variables that we've emitted a reference to.
     llvm::DenseSet<const SILGlobalVariable *> GlobalsToEmit;
 
+    // SWIFT_ENABLE_TENSORFLOW
+    /// Referenced differentiability witnesses that need to be emitted.
+    llvm::DenseSet<const SILDifferentiabilityWitness *>
+        DifferentiabilityWitnessesToEmit;
+    // SWIFT_ENABLE_TENSORFLOW END
+
     /// Additional functions we might need to serialize.
     llvm::SmallVector<const SILFunction *, 16> Worklist;
 
@@ -1061,6 +1067,7 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
   case SILInstructionKind::DifferentiabilityWitnessFunctionInst: {
     auto *dwfi = cast<DifferentiabilityWitnessFunctionInst>(&SI);
     auto *witness = dwfi->getWitness();
+    DifferentiabilityWitnessesToEmit.insert(witness);
     Mangle::ASTMangler mangler;
     auto mangledKey = mangler.mangleSILDifferentiabilityWitnessKey(
         witness->getKey());
@@ -2718,17 +2725,6 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
       writeSILDefaultWitnessTable(wt);
   }
 
-  // SWIFT_ENABLE_TENSORFLOW
-  // Write out differentiability witnesses.
-  for (const auto &diffWitness : SILMod->getDifferentiabilityWitnessList()) {
-    // TODO(TF-893): Consider checking
-    // `SILMod->shouldSerializeEntitiesAssociatedWithDeclContext` on the JVP/VJP
-    // functions.
-    if ((ShouldSerializeAll || diffWitness.isSerialized()))
-      writeSILDifferentiabilityWitness(diffWitness);
-  }
-  // SWIFT_ENABLE_TENSORFLOW END
-
   // Emit only declarations if it is a module with pre-specializations.
   // And only do it in optimized builds.
   bool emitDeclarationsForOnoneSupport =
@@ -2751,6 +2747,26 @@ void SILSerializer::writeSILBlock(const SILModule *SILMod) {
     processSILFunctionWorklist();
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  // Write out differentiability witnesses.
+  // Note: this must be done after visiting SIL functions above so that
+  // differentiability witness references (`differentiability_witness_function`
+  // instructions) have been tracked.
+  for (const auto &diffWitness : SILMod->getDifferentiabilityWitnessList()) {
+    // TODO(TF-893): Consider checking
+    // `SILMod->shouldSerializeEntitiesAssociatedWithDeclContext` on the JVP/VJP
+    // functions.
+    if ((ShouldSerializeAll || diffWitness.isSerialized()))
+      DifferentiabilityWitnessesToEmit.insert(&diffWitness);
+  }
+  for (auto *diffWitness : DifferentiabilityWitnessesToEmit)
+    writeSILDifferentiabilityWitness(*diffWitness);
+  // Process SIL functions referenced by differentiability witnesses.
+  // Note: this is necessary despite processing `FuncsToEmit` below because
+  // `Worklist` is processed separately.
+  processSILFunctionWorklist();
+  // SWIFT_ENABLE_TENSORFLOW END
+
   // Now write function declarations for every function we've
   // emitted a reference to without emitting a function body for.
   for (const SILFunction &F : *SILMod) {

test/AutoDiff/activity_analysis.swift

@@ -1,4 +1,5 @@
 // RUN: %target-swift-emit-sil -verify -Xllvm -debug-only=differentiation 2>&1 %s | %FileCheck %s
+// REQUIRES: asserts
 
 // Check that `@noDerivative` struct projections have "NONE" activity.
 
@@ -204,6 +205,106 @@ func TF_954(_ x: Float) -> Float {
 // CHECK: [ACTIVE]   %40 = begin_access [read] [static] %2 : $*Float
 // CHECK: [ACTIVE]   %41 = load [trivial] %40 : $*Float
 
+// Check `inout` argument activity.
+
+struct Mut: Differentiable {}
+extension Mut {
+  @differentiable(wrt: x)
+  mutating func mutatingMethod(_ x: Mut) -> Mut {
+    return x
+  }
+}
+
+// CHECK-LABEL: [AD] Activity info for $s17activity_analysis3MutV14mutatingMethodyA2CF at (source=0 parameters=(0))
+// CHECK: [ACTIVE] %0 = argument of bb0 : $Mut
+// CHECK: [NONE] %1 = argument of bb0 : $*Mut
+
+// TODO(TF-985): Find workaround to avoid marking non-wrt `inout` argument as
+// active.
+// expected-error @+1 {{function is not differentiable}}
+@differentiable(wrt: x)
+// expected-note @+1 {{when differentiating this function definition}}
+func nonActiveInoutArg(_ nonactive: inout Mut, _ x: Mut) -> Mut {
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  return nonactive.mutatingMethod(x)
+}
+
+// CHECK-LABEL: [AD] Activity info for $s17activity_analysis17nonActiveInoutArgyAA3MutVADz_ADtF at (source=0 parameters=(1))
+// CHECK: [ACTIVE] %0 = argument of bb0 : $*Mut
+// CHECK: [ACTIVE] %1 = argument of bb0 : $Mut
+// CHECK: [ACTIVE]   %4 = begin_access [modify] [static] %0 : $*Mut
+// CHECK: [NONE]   // function_ref Mut.mutatingMethod(_:)
+// CHECK: [ACTIVE]   %6 = apply %5(%1, %4) : $@convention(method) (Mut, @inout Mut) -> Mut
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable(wrt: x)
+// expected-note @+1 {{when differentiating this function definition}}
+func activeInoutArgMutatingMethod(_ x: Mut) -> Mut {
+  var result = x
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  result = result.mutatingMethod(result)
+  return result
+}
+
+// CHECK-LABEL: [AD] Activity info for $s17activity_analysis28activeInoutArgMutatingMethodyAA3MutVADF at (source=0 parameters=(0))
+// CHECK: [ACTIVE] %0 = argument of bb0 : $Mut
+// CHECK: [ACTIVE]   %2 = alloc_stack $Mut, var, name "result"
+// CHECK: [ACTIVE]   %4 = begin_access [read] [static] %2 : $*Mut
+// CHECK: [ACTIVE]   %5 = load [trivial] %4 : $*Mut
+// CHECK: [ACTIVE]   %7 = begin_access [modify] [static] %2 : $*Mut
+// CHECK: [NONE]   // function_ref Mut.mutatingMethod(_:)
+// CHECK: [ACTIVE]   %9 = apply %8(%5, %7) : $@convention(method) (Mut, @inout Mut) -> Mut
+// CHECK: [ACTIVE]   %11 = begin_access [modify] [static] %2 : $*Mut
+// CHECK: [ACTIVE]   %14 = begin_access [read] [static] %2 : $*Mut
+// CHECK: [ACTIVE]   %15 = load [trivial] %14 : $*Mut
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable(wrt: x)
+// expected-note @+1 {{when differentiating this function definition}}
+func activeInoutArgMutatingMethodVar(_ nonactive: inout Mut, _ x: Mut) -> Mut {
+  var result = nonactive
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  result = result.mutatingMethod(x)
+  return result
+}
+
+// CHECK_LABEL: [AD] Activity info for $s17activity_analysis31activeInoutArgMutatingMethodVaryAA3MutVADz_ADtF at (source=0 parameters=(1))
+// CHECK: [USEFUL] %0 = argument of bb0 : $*Mut
+// CHECK: [ACTIVE] %1 = argument of bb0 : $Mut
+// CHECK: [ACTIVE]   %4 = alloc_stack $Mut, var, name "result"
+// CHECK: [USEFUL]   %5 = begin_access [read] [static] %0 : $*Mut
+// CHECK: [ACTIVE]   %8 = begin_access [modify] [static] %4 : $*Mut
+// CHECK: [NONE]   // function_ref Mut.mutatingMethod(_:)
+// CHECK: [ACTIVE]   %10 = apply %9(%1, %8) : $@convention(method) (Mut, @inout Mut) -> Mut
+// CHECK: [ACTIVE]   %12 = begin_access [modify] [static] %4 : $*Mut
+// CHECK: [ACTIVE]   %15 = begin_access [read] [static] %4 : $*Mut
+// CHECK: [ACTIVE]   %16 = load [trivial] %15 : $*Mut
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable(wrt: x)
+// expected-note @+1 {{when differentiating this function definition}}
+func activeInoutArgMutatingMethodTuple(_ nonactive: inout Mut, _ x: Mut) -> Mut {
+  var result = (nonactive, x)
+  // expected-note @+1 {{cannot differentiate through 'inout' arguments}}
+  let result2 = result.0.mutatingMethod(result.0)
+  return result2
+}
+
+// CHECK-LABEL: [AD] Activity info for $s17activity_analysis33activeInoutArgMutatingMethodTupleyAA3MutVADz_ADtF at (source=0 parameters=(1))
+// CHECK: [USEFUL] %0 = argument of bb0 : $*Mut
+// CHECK: [ACTIVE] %1 = argument of bb0 : $Mut
+// CHECK: [ACTIVE]   %4 = alloc_stack $(Mut, Mut), var, name "result"
+// CHECK: [ACTIVE]   %5 = tuple_element_addr %4 : $*(Mut, Mut), 0
+// CHECK: [ACTIVE]   %6 = tuple_element_addr %4 : $*(Mut, Mut), 1
+// CHECK: [USEFUL]   %7 = begin_access [read] [static] %0 : $*Mut
+// CHECK: [ACTIVE]   %11 = begin_access [read] [static] %4 : $*(Mut, Mut)
+// CHECK: [ACTIVE]   %12 = tuple_element_addr %11 : $*(Mut, Mut), 0
+// CHECK: [ACTIVE]   %13 = load [trivial] %12 : $*Mut
+// CHECK: [ACTIVE]   %15 = begin_access [modify] [static] %4 : $*(Mut, Mut)
+// CHECK: [ACTIVE]   %16 = tuple_element_addr %15 : $*(Mut, Mut), 0
+// CHECK: [NONE]   // function_ref Mut.mutatingMethod(_:)
+// CHECK: [ACTIVE]   %18 = apply %17(%13, %16) : $@convention(method) (Mut, @inout Mut) -> Mut
+
 //===----------------------------------------------------------------------===//
 // Non-differentiable functions
 //===----------------------------------------------------------------------===//