[AutoDiff] Propagate '@nondiff' from AST function types to SIL function types.

include/swift/AST/DiagnosticsSIL.def

@@ -396,6 +396,8 @@ NOTE(autodiff_protocol_member_not_differentiable,none,
 NOTE(autodiff_protocol_member_subset_indices_not_differentiable,none,
      "member is differentiable only with respect to a smaller subset of "
      "arguments", ())
+NOTE(autodiff_function_nondiff_parameter_not_differentiable,none,
+     "cannot differentiate with respect to a '@nondiff' parameter", ())
 NOTE(autodiff_function_assoc_func_requirements_unmet,none,
      "function call is not differentiable because generic requirements are not "
      "met", ())

include/swift/AST/Types.h

@@ -3093,6 +3093,7 @@ class AnyFunctionType : public TypeBase {
     return getExtInfo().getRepresentation();
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
   /// Given `indices`, `differentiationOrder`, and `kind`, calculates the type
   /// of the corresponding autodiff associated function.
   ///
@@ -3105,6 +3106,8 @@ class AnyFunctionType : public TypeBase {
       LookupConformanceFn lookupConformance,
       GenericSignature *whereClauseGenericSignature = nullptr);
 
+  AnyFunctionType *getWithoutDifferentiability() const;
+
   /// \brief True if the parameter declaration it is attached to is guaranteed
   /// to not persist the closure for longer than the duration of the call.
   bool isNoEscape() const {

include/swift/SIL/SILBuilder.h

@@ -547,7 +547,6 @@ class SILBuilder {
         theFunction));
   }
 
-
   AutoDiffFunctionExtractInst *createAutoDiffFunctionExtractOriginal(
       SILLocation loc, SILValue theFunction) {
     return insert(new (getModule()) AutoDiffFunctionExtractInst(

lib/AST/ASTPrinter.cpp

@@ -2338,7 +2338,7 @@ static void printParameterFlags(ASTPrinter &printer, PrintOptions options,
     printer << "@escaping ";
   // SWIFT_ENABLE_TENSORFLOW
   if (!options.excludeAttrKind(TAK_nondiff) && flags.isNonDifferentiable())
-    printer << "@nondiff";
+    printer << "@nondiff ";
 
   switch (flags.getValueOwnership()) {
   case ValueOwnership::Default:

lib/AST/AutoDiff.cpp

@@ -254,15 +254,12 @@ AutoDiffParameterIndices::getLowered(AnyFunctionType *functionType) const {
 }
 
 static unsigned getNumAutoDiffParameterIndices(AnyFunctionType *fnTy) {
-  unsigned numAutoDiffParameterIndices = 0;
-  // FIXME: Compute the exact parameter count.
-  // Do not loop ad-infinitum; loop either 1 or 2 iterations, depending on
-  // whether the function is a free function/static method/instance method.
-  while (fnTy != nullptr) {
-    numAutoDiffParameterIndices += fnTy->getNumParams();
-    fnTy = fnTy->getResult()->getAs<AnyFunctionType>();
-  }
-  return numAutoDiffParameterIndices;
+  // TODO: For more correct counting, we still need to know whether it's a
+  // method or not.
+  unsigned numParameters = fnTy->getNumParams();
+  if (auto *innerFn = fnTy->getResult()->getAs<AnyFunctionType>())
+    numParameters += innerFn->getNumParams();
+  return numParameters;
 }
 
 /// Returns true if the given type conforms to `Differentiable` in the given

lib/AST/Builtins.cpp

@@ -1075,8 +1075,8 @@ static ValueDecl *getAutoDiffApplyAssociatedFunction(
   // of the associated function type.
   auto *origFnTy =
       firstArgGen.build(builder)->castTo<AnyFunctionType>();
-  origFnTy = origFnTy->withExtInfo(
-      origFnTy->getExtInfo().withDifferentiable(false).withNoEscape(false));
+  origFnTy = origFnTy->getWithoutDifferentiability()->withExtInfo(
+      origFnTy->getExtInfo().withNoEscape(false));
   auto autodiffBuilder = AutoDiffParameterIndicesBuilder::inferParameters(
       origFnTy, Context.getStdlibModule());
   auto *paramIndices = autodiffBuilder.build(Context);

lib/AST/Type.cpp

@@ -4304,3 +4304,18 @@ AnyFunctionType *AnyFunctionType::getAutoDiffAssociatedFunctionType(
 
   return associatedFunction;
 }
+
+AnyFunctionType *AnyFunctionType::getWithoutDifferentiability() const {
+  SmallVector<Param, 8> newParams;
+  for (auto &param : getParams()) {
+    Param newParam(param.getPlainType(), param.getLabel(),
+                   param.getParameterFlags().withNonDifferentiable(false));
+    newParams.push_back(newParam);
+  }
+  auto nonDiffExtInfo = getExtInfo().withDifferentiable(false);
+  if (isa<FunctionType>(this))
+    return FunctionType::get(newParams, getResult(), nonDiffExtInfo);
+  assert(isa<GenericFunctionType>(this));
+  return GenericFunctionType::get(getOptGenericSignature(), newParams,
+                                  getResult(), nonDiffExtInfo);
+}

lib/IRGen/GenDiffFunc.cpp

@@ -38,12 +38,16 @@ using DiffFuncIndex =
 namespace {
 class DiffFuncFieldInfo final : public RecordField<DiffFuncFieldInfo> {
 public:
-  DiffFuncFieldInfo(DiffFuncIndex index, const TypeInfo &type)
-      : RecordField(type), Index(index) {}
+  DiffFuncFieldInfo(DiffFuncIndex index, const TypeInfo &type,
+                    const SmallBitVector &parameterIndices)
+      : RecordField(type), Index(index), ParameterIndices(parameterIndices) {}
 
   /// The field index.
   const DiffFuncIndex Index;
 
+  /// The parameter indices.
+  SmallBitVector ParameterIndices;
+
   std::string getFieldName() const {
     auto extractee = std::get<0>(Index);
     auto differentiationOrder = std::get<1>(Index);
@@ -59,17 +63,14 @@ class DiffFuncFieldInfo final : public RecordField<DiffFuncFieldInfo> {
 
   SILType getType(IRGenModule &IGM, SILType t) const {
     auto fnTy = t.castTo<SILFunctionType>();
-    auto extInfo = fnTy->getExtInfo();
-    auto nondiffExtInfo = extInfo.withDifferentiable(false);
-    auto origFnTy = fnTy->getWithExtInfo(nondiffExtInfo);
+    auto origFnTy = fnTy->getWithoutDifferentiability();
     if (std::get<0>(Index) == AutoDiffFunctionExtractInst::Extractee::Original)
       return SILType::getPrimitiveObjectType(origFnTy);
     auto differentiationOrder = std::get<1>(Index);
     auto kind = *std::get<0>(Index).getExtracteeAsAssociatedFunction();
     auto assocTy = origFnTy->getAutoDiffAssociatedFunctionType(
-        SmallBitVector(origFnTy->getNumParameters(), true), /*resultIndex*/ 0,
-        differentiationOrder, kind, IGM.getSILModule(),
-        LookUpConformanceInModule(IGM.getSwiftModule()));
+        ParameterIndices, /*resultIndex*/ 0, differentiationOrder, kind,
+        IGM.getSILModule(), LookUpConformanceInModule(IGM.getSwiftModule()));
     return SILType::getPrimitiveObjectType(assocTy);
   }
 };
@@ -118,14 +119,13 @@ class DiffFuncTypeBuilder
                                DiffFuncIndex> {
 
   SILFunctionType *origFnTy;
+  SmallBitVector parameterIndices;
 
 public:
   DiffFuncTypeBuilder(IRGenModule &IGM, SILFunctionType *fnTy)
-      : RecordTypeBuilder(IGM) {
+      : RecordTypeBuilder(IGM), origFnTy(fnTy->getWithoutDifferentiability()),
+        parameterIndices(fnTy->getDifferentiationParameterIndices()) {
     assert(fnTy->isDifferentiable());
-    auto extInfo = fnTy->getExtInfo();
-    auto nondiffExtInfo = extInfo.withDifferentiable(false);
-    origFnTy = fnTy->getWithExtInfo(nondiffExtInfo);
   }
 
   TypeInfo *createFixed(ArrayRef<DiffFuncFieldInfo> fields,
@@ -150,7 +150,7 @@ class DiffFuncTypeBuilder
 
   DiffFuncFieldInfo getFieldInfo(unsigned index, DiffFuncIndex field,
                                  const TypeInfo &fieldTI) {
-    return DiffFuncFieldInfo(field, fieldTI);
+    return DiffFuncFieldInfo(field, fieldTI, parameterIndices);
   }
 
   SILType getType(DiffFuncIndex field) {
@@ -159,9 +159,8 @@ class DiffFuncTypeBuilder
     auto differentiationOrder = std::get<1>(field);
     auto kind = *std::get<0>(field).getExtracteeAsAssociatedFunction();
     auto assocTy = origFnTy->getAutoDiffAssociatedFunctionType(
-        SmallBitVector(origFnTy->getNumParameters(), true), /*resultIndex*/ 0,
-        differentiationOrder, kind, IGM.getSILModule(),
-        LookUpConformanceInModule(IGM.getSwiftModule()));
+        parameterIndices, /*resultIndex*/ 0, differentiationOrder, kind,
+        IGM.getSILModule(), LookUpConformanceInModule(IGM.getSwiftModule()));
     return SILType::getPrimitiveObjectType(assocTy);
   }
 

lib/IRGen/IRGenSIL.cpp

@@ -3580,10 +3580,8 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
     (void)adFnExp.claimAll();
     tmpCalleeLV = LoweredValue(e);
 
-    origCalleeType = origCalleeType->getWithExtInfo(
-        origCalleeType->getExtInfo().withDifferentiable(false));
-    substCalleeType = substCalleeType->getWithExtInfo(
-        substCalleeType->getExtInfo().withDifferentiable(false));
+    origCalleeType = origCalleeType->getWithoutDifferentiability();
+    substCalleeType = substCalleeType->getWithoutDifferentiability();
   }
   const LoweredValue &calleeLV =
       tmpCalleeLV ? *tmpCalleeLV : getLoweredValue(site.getCallee());