[AutoDiff] Add Type Checking for @transposing (#25684)

Do the type checking for top level functions and methods that are marked as the transpose of another function.

Author: bartchr808

Diff for: include/swift/AST/Attr.def

@@ -432,6 +432,9 @@ SIMPLE_DECL_ATTR(compilerEvaluable, CompilerEvaluable,
   /* Not serialized */ 91)
 SIMPLE_DECL_ATTR(noDerivative, NoDerivative,
   OnVar, 92)
+DECL_ATTR(transposing, Transposing,
+  OnFunc | LongAttribute | AllowMultipleAttributes |
+  NotSerialized, 93)
 
 #undef TYPE_ATTR
 #undef DECL_ATTR_ALIAS

Diff for: include/swift/AST/Attr.h

@@ -1682,6 +1682,79 @@ class DifferentiatingAttr final
     return DA->getKind() == DAK_Differentiating;
   }
 };
+  
+/// Attribute that registers a function as a transpose of another function.
+///
+/// Examples:
+///   @transposing(foo)
+///   @transposing(+, wrt: (lhs, rhs))
+class TransposingAttr final
+      : public DeclAttribute,
+        private llvm::TrailingObjects<DifferentiableAttr,
+                                      ParsedAutoDiffParameter> {
+  /// The base type of the original function.
+  /// This is non-null only when the original function is not top-level (i.e. it
+  /// is an instance/static method).
+  TypeRepr *BaseType;
+  /// The original function name.
+  DeclNameWithLoc Original;
+  /// The original function, resolved by the type checker.
+  FuncDecl *OriginalFunction = nullptr;
+  /// The number of parsed parameters specified in 'wrt:'.
+  unsigned NumParsedParameters = 0;
+  /// The differentiation parameters' indices, resolved by the type checker.
+  AutoDiffIndexSubset *ParameterIndexSubset = nullptr;
+  
+  explicit TransposingAttr(ASTContext &context, bool implicit,
+                           SourceLoc atLoc, SourceRange baseRange,
+                           TypeRepr *baseType, DeclNameWithLoc original,
+                           ArrayRef<ParsedAutoDiffParameter> params);
+  
+  explicit TransposingAttr(ASTContext &context, bool implicit,
+                           SourceLoc atLoc, SourceRange baseRange,
+                           TypeRepr *baseType, DeclNameWithLoc original,
+                           AutoDiffIndexSubset *indices);
+  
+public:
+  static TransposingAttr *create(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 ArrayRef<ParsedAutoDiffParameter> params);
+  
+  static TransposingAttr *create(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 AutoDiffIndexSubset *indices);
+  
+  TypeRepr *getBaseType() const { return BaseType; }
+  DeclNameWithLoc getOriginal() const { return Original; }
+  
+  FuncDecl *getOriginalFunction() const { return OriginalFunction; }
+  void setOriginalFunction(FuncDecl *decl) { OriginalFunction = decl; }
+  
+  /// The parsed transposing parameters, i.e. the list of parameters
+  /// specified in 'wrt:'.
+  ArrayRef<ParsedAutoDiffParameter> getParsedParameters() const {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  MutableArrayRef<ParsedAutoDiffParameter> getParsedParameters() {
+    return {getTrailingObjects<ParsedAutoDiffParameter>(), NumParsedParameters};
+  }
+  size_t numTrailingObjects(OverloadToken<ParsedAutoDiffParameter>) const {
+    return NumParsedParameters;
+  }
+  
+  AutoDiffIndexSubset *getParameterIndexSubset() const {
+    return ParameterIndexSubset;
+  }
+  void setParameterIndices(AutoDiffIndexSubset *pi) {
+    ParameterIndexSubset = pi;
+  }
+  
+  static bool classof(const DeclAttribute *DA) {
+    return DA->getKind() == DAK_Transposing;
+  }
+};
 
 /// Attributes that may be applied to declarations.
 class DeclAttributes {

Diff for: include/swift/AST/AutoDiff.h

@@ -354,6 +354,13 @@ class AutoDiffIndexSubset : public llvm::FoldingSetNode {
   unsigned getNumIndices() const {
     return (unsigned)std::distance(begin(), end());
   }
+  
+  SmallBitVector getBitVector() const {
+    SmallBitVector indicesBitVec(capacity, false);
+    for (auto index : getIndices())
+      indicesBitVec.set(index);
+    return indicesBitVec;
+  }
 
   bool contains(unsigned index) const {
     unsigned bitWordIndex, offset;

Diff for: include/swift/AST/DiagnosticsParse.def

@@ -1512,13 +1512,23 @@ ERROR(attr_differentiating_expected_original_name,PointsToFirstBadToken,
 ERROR(attr_differentiating_expected_label_linear_or_wrt,none,
       "expected either 'linear' or 'wrt:'", ())
 
+// transposing
+ERROR(attr_transposing_expected_original_name,PointsToFirstBadToken,
+      "expected an original function name", ())
+ERROR(attr_transposing_expected_label_linear_or_wrt,none,
+      "expected 'wrt:'", ())
+
+// transposing `wrt` parameters clause
+ERROR(transposing_params_clause_expected_parameter,PointsToFirstBadToken,
+      "expected a parameter, which can be a 'unsigned int' parameter number "
+      "or 'self'", ())
+
 // differentiation `wrt` parameters clause
 ERROR(expected_colon_after_label,PointsToFirstBadToken,
       "expected a colon ':' after '%0'", (StringRef))
 ERROR(diff_params_clause_expected_parameter,PointsToFirstBadToken,
       "expected a parameter, which can be a function parameter name, "
-      "parameter index, or 'self'",
-      ())
+      "parameter index, or 'self'", ())
 
 // [differentiable ...] (sil-decl attr)
 ERROR(sil_attr_differentiable_expected_keyword,PointsToFirstBadToken,

Diff for: include/swift/AST/DiagnosticsSema.def

@@ -2826,6 +2826,19 @@ ERROR(differentiating_attr_not_in_same_file_as_original,none,
 ERROR(differentiating_attr_original_already_has_derivative,none,
       "a derivative already exists for %0", (DeclName))
 
+// transposing
+ERROR(transpose_params_clause_param_not_differentiable,none,
+      "can only transpose with respect to parameters that conform to "
+      "'Differentiable' and where '%0 == %0.TangentVector'", (StringRef))
+ERROR(transposing_attr_overload_not_found,none,
+      "could not find function %0 with expected type %1", (DeclName, Type))
+ERROR(transposing_attr_cant_use_named_wrt_params,none,
+      "cannot use named wrt parameters in '@transposing' attribute, found %0",
+      (Identifier))
+ERROR(transposing_attr_result_value_not_differentiable,none,
+      "'@transposing' attribute requires original function result to "
+      "conform to 'Differentiable'", (Type))
+
 // differentiation `wrt` parameters clause
 ERROR(diff_function_no_parameters,none,
       "%0 has no parameters to differentiate with respect to", (DeclName))

Diff for: include/swift/AST/Types.h

@@ -19,6 +19,7 @@
 
 // SWIFT_ENABLE_TENSORFLOW
 #include "swift/AST/AutoDiff.h"
+#include "swift/AST/Attr.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/GenericParamKey.h"
 #include "swift/AST/Identifier.h"
@@ -3093,7 +3094,7 @@ class AnyFunctionType : public TypeBase {
   ///
   /// If `makeSelfParamFirst` is true, self's tangent is reordered to appear
   /// first. This should be used during type-checking, e.g. type-checking
-  /// `@differentiable` and `@differentiating` attributes.
+  /// `@differentiable`, `@differentiating`, and `@transposing` attributes.
   ///
   /// \note The original function type (`self`) need not be `@differentiable`.
   /// The resulting function will preserve all `ExtInfo` of the original
@@ -3108,6 +3109,13 @@ class AnyFunctionType : public TypeBase {
   /// Given the type of an autodiff associated function, returns the
   /// corresponding original function type.
   AnyFunctionType *getAutoDiffOriginalFunctionType();
+  
+  /// Given the type of a transposing associated function, returns the
+  /// corresponding original function type.
+  AnyFunctionType *
+  getTransposeOriginalFunctionType(TransposingAttr *attr,
+                                   AutoDiffIndexSubset *wrtParamIndices,
+                                   bool wrtSelf);
 
   AnyFunctionType *getWithoutDifferentiability() const;
 

Diff for: include/swift/Parse/Parser.h

@@ -960,10 +960,17 @@ class Parser {
   /// Parse a differentiation parameters clause.
   bool parseDifferentiationParametersClause(
       SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName);
+  
+  /// Parse a transposing parameters clause.
+  bool parseTransposingParametersClause(
+      SmallVectorImpl<ParsedAutoDiffParameter> &params, StringRef attrName);
 
   /// Parse the @differentiating attribute.
   ParserResult<DifferentiatingAttr>
   parseDifferentiatingAttribute(SourceLoc AtLoc, SourceLoc Loc);
+  
+  ParserResult<TransposingAttr> parseTransposingAttribute(SourceLoc AtLoc,
+                                                          SourceLoc Loc);
 
   /// Parse a specific attribute.
   ParserStatus parseDeclAttribute(DeclAttributes &Attributes, SourceLoc AtLoc);

Diff for: include/swift/Serialization/ModuleFormat.h

@@ -1670,6 +1670,15 @@ namespace decls_block {
     DeclIDField, // Original function declaration.
     BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.
   >;
+  
+  // SWIFT_ENABLE_TENSORFLOW
+  using TransposingDeclAttrLayout = BCRecordLayout<
+    Transposing_DECL_ATTR,
+    BCFixed<1>, // Implicit flag.
+    IdentifierIDField, // Original name.
+    DeclIDField, // Original function declaration.
+    BCArray<BCFixed<1>> // Differentiation parameter indices' bitvector.
+  >;
 
 #define SIMPLE_DECL_ATTR(X, CLASS, ...) \
   using CLASS##DeclAttrLayout = BCRecordLayout< \

Diff for: lib/AST/Attr.cpp

@@ -389,6 +389,67 @@ static std::string getDifferentiationParametersClauseString(
   return printer.str();
 }
 
+// Returns the differentiation parameters clause string for the given function,
+// parameter indices, and parsed parameters.
+static std::string getTransposingParametersClauseString(
+    const AbstractFunctionDecl *function, AutoDiffIndexSubset *indices,
+    ArrayRef<ParsedAutoDiffParameter> parsedParams) {
+  bool isInstanceMethod = function && function->isInstanceMember();
+  
+  std::string result;
+  llvm::raw_string_ostream printer(result);
+  
+  // Use parameters from `AutoDiffIndexSubset`, if specified.
+  if (indices) {
+    SmallBitVector parameters(indices->getBitVector());
+    auto parameterCount = parameters.count();
+    printer << "wrt: ";
+    if (parameterCount > 1)
+      printer << '(';
+    // Check if differentiating wrt `self`. If so, manually print it first.
+    if (isInstanceMethod && parameters.test(parameters.size() - 1)) {
+      parameters.reset(parameters.size() - 1);
+      printer << "self";
+      if (parameters.any())
+        printer << ", ";
+    }
+    // Print remaining differentiation parameters.
+    interleave(parameters.set_bits(), [&](unsigned index) { printer << index; },
+               [&] { printer << ", "; });
+    if (parameterCount > 1)
+      printer << ')';
+  }
+  // Otherwise, use the parsed parameters.
+  else if (!parsedParams.empty()) {
+    printer << "wrt: ";
+    if (parsedParams.size() > 1)
+      printer << '(';
+    interleave(
+        parsedParams,
+        [&](const ParsedAutoDiffParameter &param) {
+          switch (param.getKind()) {
+            case ParsedAutoDiffParameter::Kind::Named:
+              printer << param.getName();
+              break;
+            case ParsedAutoDiffParameter::Kind::Self:
+              printer << "self";
+              break;
+            case ParsedAutoDiffParameter::Kind::Ordered:
+              assert((param.getIndex() < function->getParameters()->size()) &&
+                      "'wrt:' parameter index should be less than the number "
+                      "of parameters");
+            auto *funcParam = function->getParameters()->get(param.getIndex());
+            printer << funcParam->getNameStr();
+            break;
+          }
+        },
+        [&] { printer << ", "; });
+    if (parsedParams.size() > 1)
+      printer << ')';
+  }
+  return printer.str();
+}
+
 // SWIFT_ENABLE_TENSORFLOW
 // Print the arguments of the given `@differentiable` attribute.
 static void printDifferentiableAttrArguments(
@@ -803,6 +864,22 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
     Printer << ')';
     break;
   }
+      
+  // SWIFT_ENABLE_TENSORFLOW
+  case DAK_Transposing: {
+    Printer.printAttrName("@transposing");
+    Printer << '(';
+    auto *attr = cast<TransposingAttr>(this);
+    auto *transpose = dyn_cast_or_null<AbstractFunctionDecl>(D);
+    Printer << attr->getOriginal().Name;
+    auto diffParamsString = getTransposingParametersClauseString(
+        transpose, attr->getParameterIndexSubset(),
+        attr->getParsedParameters());
+    if (!diffParamsString.empty())
+      Printer << ", " << diffParamsString;
+    Printer << ')';
+    break;
+  }
 
   case DAK_DynamicReplacement: {
     Printer.printAttrName("@_dynamicReplacement");
@@ -955,6 +1032,8 @@ StringRef DeclAttribute::getAttrName() const {
     return "differentiable";
   case DAK_Differentiating:
     return "differentiating";
+  case DAK_Transposing:
+    return "transposing";
   }
   llvm_unreachable("bad DeclAttrKind");
 }
@@ -1438,9 +1517,49 @@ DifferentiatingAttr::create(ASTContext &context, bool implicit,
                                        std::move(original), linear, indices);
 }
 
+TransposingAttr::TransposingAttr(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 ArrayRef<ParsedAutoDiffParameter> params)
+    : DeclAttribute(DAK_Transposing, atLoc, baseRange, implicit),
+      BaseType(baseType), Original(std::move(original)),
+      NumParsedParameters(params.size()) {
+  std::uninitialized_copy(params.begin(), params.end(),
+                          getTrailingObjects<ParsedAutoDiffParameter>());
+}
+
+TransposingAttr::TransposingAttr(ASTContext &context, bool implicit,
+                                 SourceLoc atLoc, SourceRange baseRange,
+                                 TypeRepr *baseType, DeclNameWithLoc original,
+                                 AutoDiffIndexSubset *indices)
+    : DeclAttribute(DAK_Differentiating, atLoc, baseRange, implicit),
+      BaseType(baseType), Original(std::move(original)),
+      ParameterIndexSubset(indices) {}
+
+TransposingAttr *
+TransposingAttr::create(ASTContext &context, bool implicit, SourceLoc atLoc,
+                        SourceRange baseRange, TypeRepr *baseType,
+                        DeclNameWithLoc original,
+                        ArrayRef<ParsedAutoDiffParameter> params) {
+  unsigned size = totalSizeToAlloc<ParsedAutoDiffParameter>(params.size());
+  void *mem = context.Allocate(size, alignof(TransposingAttr));
+  return new (mem) TransposingAttr(context, implicit, atLoc, baseRange,
+                                   baseType, std::move(original), params);
+}
+
+TransposingAttr *
+TransposingAttr::create(ASTContext &context, bool implicit, SourceLoc atLoc,
+                        SourceRange baseRange, TypeRepr *baseType,
+                        DeclNameWithLoc original,
+                        AutoDiffIndexSubset *indices) {
+  void *mem =
+      context.Allocate(sizeof(TransposingAttr), alignof(TransposingAttr));
+  return new (mem) TransposingAttr(context, implicit, atLoc, baseRange,
+                                   baseType, std::move(original), indices);
+}
+
 ImplementsAttr::ImplementsAttr(SourceLoc atLoc, SourceRange range,
-                               TypeLoc ProtocolType,
-                               DeclName MemberName,
+                               TypeLoc ProtocolType, DeclName MemberName,
                                DeclNameLoc MemberNameLoc)
     : DeclAttribute(DAK_Implements, atLoc, range, /*Implicit=*/false),
       ProtocolType(ProtocolType),