[AutoDiff upstream] Upstream serialization changes. (#30720)

Upstream random serialization code from tensorflow branch:
- `SerializedSILLoader::getAllDifferentiabilityWitnesses()`
- Add diff. witness serialization for functions with `inout` parameters.
- `deserializeSILFunctionType`: fix assertion for differentiability kind.

Author: dan-zheng

include/swift/Serialization/SerializedSILLoader.h

@@ -13,6 +13,7 @@
 #ifndef SWIFT_SERIALIZATION_SILLOADER_H
 #define SWIFT_SERIALIZATION_SILLOADER_H
 
+#include "swift/AST/AutoDiff.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Identifier.h"
 #include "swift/SIL/Notifications.h"
@@ -102,6 +103,9 @@ class SerializedSILLoader {
   /// Deserialize all Properties in all SILModules.
   void getAllProperties();
 
+  /// Deserialize all DifferentiabilityWitnesses in all SILModules.
+  void getAllDifferentiabilityWitnesses();
+
   SerializedSILLoader(const SerializedSILLoader &) = delete;
   SerializedSILLoader(SerializedSILLoader &&) = delete;
   SerializedSILLoader &operator=(const SerializedSILLoader &) = delete;

lib/Serialization/Deserialization.cpp

@@ -5411,8 +5411,11 @@ class TypeDeserializer {
     };
 
     // Bounds check.  FIXME: overflow
-    if (2 * numParams + 2 * numResults + 2 * unsigned(hasErrorResult)
-          > variableData.size()) {
+    unsigned entriesPerParam =
+        diffKind != DifferentiabilityKind::NonDifferentiable ? 3 : 2;
+    if (entriesPerParam * numParams + 2 * numResults +
+            2 * unsigned(hasErrorResult) >
+        variableData.size()) {
       MF.fatal();
     }
 

lib/Serialization/DeserializeSIL.cpp

@@ -3573,19 +3573,22 @@ SILDeserializer::readDifferentiabilityWitness(DeclID DId) {
   }
   auto derivativeGenSig = MF->getGenericSignature(derivativeGenSigID);
 
+  auto originalFnType = original->getLoweredFunctionType();
   SmallVector<unsigned, 8> parameterAndResultIndices(
       rawParameterAndResultIndices.begin(), rawParameterAndResultIndices.end());
   assert(parameterAndResultIndices.size() ==
              numParameterIndices + numResultIndices &&
          "Parameter/result indices count mismatch");
-  auto *parameterIndices = IndexSubset::get(
-      MF->getContext(), original->getLoweredFunctionType()->getNumParameters(),
-      ArrayRef<unsigned>(parameterAndResultIndices)
-          .take_front(numParameterIndices));
-  auto *resultIndices = IndexSubset::get(
-      MF->getContext(), original->getLoweredFunctionType()->getNumResults(),
-      ArrayRef<unsigned>(parameterAndResultIndices)
-          .take_back(numResultIndices));
+  auto *parameterIndices =
+      IndexSubset::get(MF->getContext(), originalFnType->getNumParameters(),
+                       ArrayRef<unsigned>(parameterAndResultIndices)
+                           .take_front(numParameterIndices));
+  auto numResults = originalFnType->getNumResults() +
+                    originalFnType->getNumIndirectMutatingParameters();
+  auto *resultIndices =
+      IndexSubset::get(MF->getContext(), numResults,
+                       ArrayRef<unsigned>(parameterAndResultIndices)
+                           .take_back(numResultIndices));
 
   AutoDiffConfig config(parameterIndices, resultIndices, derivativeGenSig);
   auto *diffWitness =

lib/Serialization/Serialization.cpp

@@ -2438,11 +2438,6 @@ class Serializer::DeclSerializer : public DeclVisitor<DeclSerializer> {
     case DAK_Transpose: {
       auto abbrCode = S.DeclTypeAbbrCodes[TransposeDeclAttrLayout::Code];
       auto *attr = cast<TransposeAttr>(DA);
-      // NOTE(TF-838): `@transpose` attribute serialization is blocked by
-      // `@transpose` attribute type-checking (TF-830), which resolves
-      // the original declaration.
-      if (!attr->getOriginalFunction())
-        return;
       assert(attr->getOriginalFunction() &&
              "`@transpose` attribute should have original declaration set "
              "during construction or parsing");

lib/Serialization/SerializeSIL.cpp

@@ -2542,7 +2542,11 @@ void SILSerializer::writeSILDifferentiabilityWitness(
              dw.getParameterIndices()->getCapacity() &&
          "Original function parameter count should match differentiability "
          "witness parameter indices capacity");
-  assert(originalFnType->getNumResults() ==
+  unsigned numInoutParameters = llvm::count_if(
+      originalFnType->getParameters(), [](SILParameterInfo paramInfo) {
+        return paramInfo.isIndirectMutating();
+      });
+  assert(originalFnType->getNumResults() + numInoutParameters ==
              dw.getResultIndices()->getCapacity() &&
          "Original function result count should match differentiability "
          "witness result indices capacity");

lib/Serialization/SerializedSILLoader.cpp

@@ -202,3 +202,8 @@ void SerializedSILLoader::getAllProperties() {
     Des->getAllProperties();
 }
 
+/// Deserialize all DifferentiabilityWitnesses in all SILModules.
+void SerializedSILLoader::getAllDifferentiabilityWitnesses() {
+  for (auto &Des : LoadedSILSections)
+    Des->getAllDifferentiabilityWitnesses();
+}