Merge pull request #4455 from swiftwasm/main

[pull] swiftwasm from main

Author: kateinoigakukun

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -549,6 +549,10 @@ class PullbackCloner::Implementation final
     if (auto adjProj = getAdjointProjection(origBB, originalValue))
       return (bufferMap[{origBB, originalValue}] = adjProj);
 
+    LLVM_DEBUG(getADDebugStream() << "Creating new adjoint buffer for"
+               << originalValue
+               << "in bb" << origBB->getDebugID() << '\n');
+
     auto bufType = getRemappedTangentType(originalValue->getType());
     // Set insertion point for local allocation builder: before the last local
     // allocation, or at the start of the pullback function's entry if no local
@@ -583,6 +587,12 @@ class PullbackCloner::Implementation final
     assert(originalValue->getFunction() == &getOriginal());
     assert(rhsAddress->getFunction() == &getPullback());
     auto adjointBuffer = getAdjointBuffer(origBB, originalValue);
+
+    LLVM_DEBUG(getADDebugStream() << "Adding"
+               << rhsAddress << "to adjoint of "
+               << originalValue
+               << "in bb" << origBB->getDebugID() << '\n');
+
     builder.emitInPlaceAdd(loc, adjointBuffer, rhsAddress);
   }
 
@@ -2340,7 +2350,9 @@ SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
   for (auto activeValue : predBBActiveValues) {
     LLVM_DEBUG(getADDebugStream()
                << "Propagating adjoint of active value " << activeValue
-               << " to predecessors' pullback blocks\n");
+               << "from bb" << origBB->getDebugID()
+               << " to predecessors' (bb" << origPredBB->getDebugID()
+               << ") pullback blocks\n");
     switch (getTangentValueCategory(activeValue)) {
     case SILValueCategory::Object: {
       auto activeValueAdj = getAdjointValue(origBB, activeValue);
@@ -2511,14 +2523,13 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
     case SILValueCategory::Address: {
       auto bbArgAdjBuf = getAdjointBuffer(bb, bbArg);
       for (auto pair : incomingValues) {
-        auto *predBB = std::get<0>(pair);
         auto incomingValue = std::get<1>(pair);
         // Handle `switch_enum` on `Optional`.
         auto termInst = bbArg->getSingleTerminator();
         if (isSwitchEnumInstOnOptional(termInst))
           accumulateAdjointForOptional(bb, incomingValue, bbArgAdjBuf);
         else
-          addToAdjointBuffer(predBB, incomingValue, bbArgAdjBuf, pbLoc);
+          addToAdjointBuffer(bb, incomingValue, bbArgAdjBuf, pbLoc);
       }
       break;
     }

lib/Sema/CSDiagnostics.cpp

@@ -3970,8 +3970,11 @@ bool InvalidMemberRefOnExistential::diagnoseAsError() {
 
   // If the base expression is a reference to a function or subscript
   // parameter, offer a fixit that replaces the existential parameter type with
-  // its generic equivalent, e.g. func foo(p: any P) → func foo<T: P>(p: T).
-  // FIXME: Add an option to use 'some' vs. an explicit generic parameter.
+  // its generic equivalent, e.g. func foo(p: any P) → func foo(p: some P).
+  // Replacing 'any' with 'some' allows the code to compile without further
+  // changes, such as naming an explicit type parameter, and is future-proofed
+  // for same-type requirements on primary associated types instead of needing
+  // a where clause.
 
   if (!PD || !PD->getDeclContext()->getAsDecl())
     return true;
@@ -4020,70 +4023,42 @@ bool InvalidMemberRefOnExistential::diagnoseAsError() {
   if (PD->isInOut())
     return true;
 
-  constexpr StringRef GPNamePlaceholder = "<#generic parameter name#>";
-  SourceRange TyReplacementRange;
-  SourceRange RemoveAnyRange;
-  SourceLoc GPDeclLoc;
-  std::string GPDeclStr;
-  {
-    llvm::raw_string_ostream OS(GPDeclStr);
-    auto *const GC = PD->getDeclContext()->getAsDecl()->getAsGenericContext();
-    if (GC->getParsedGenericParams()) {
-      GPDeclLoc = GC->getParsedGenericParams()->getRAngleLoc();
-      OS << ", ";
-    } else {
-      GPDeclLoc =
-          isa<AbstractFunctionDecl>(GC)
-              ? cast<AbstractFunctionDecl>(GC)->getParameters()->getLParenLoc()
-              : cast<SubscriptDecl>(GC)->getIndices()->getLParenLoc();
-      OS << "<";
-    }
-    OS << GPNamePlaceholder << ": ";
-
-    auto *TR = PD->getTypeRepr()->getWithoutParens();
-    if (auto *STR = dyn_cast<SpecifierTypeRepr>(TR)) {
-      TR = STR->getBase()->getWithoutParens();
-    }
-    if (auto *ETR = dyn_cast<ExistentialTypeRepr>(TR)) {
-      TR = ETR->getConstraint();
-      RemoveAnyRange = SourceRange(ETR->getAnyLoc(), TR->getStartLoc());
-      TR = TR->getWithoutParens();
-    }
-    if (auto *MTR = dyn_cast<MetatypeTypeRepr>(TR)) {
-      TR = MTR->getBase();
-
-      // (P & Q).Type -> T.Type
-      // (P).Type -> (T).Type
-      // ((P & Q)).Type -> ((T)).Type
-      if (auto *TTR = dyn_cast<TupleTypeRepr>(TR)) {
-        assert(TTR->isParenType());
-        if (!isa<CompositionTypeRepr>(TTR->getElementType(0))) {
-          TR = TR->getWithoutParens();
-        }
-      }
-    }
-    TyReplacementRange = TR->getSourceRange();
-
-    // Strip any remaining parentheses and print the conformance constraint.
-    TR->getWithoutParens()->print(OS);
-
-    if (!GC->getParsedGenericParams()) {
-      OS << ">";
-    }
+  auto *typeRepr = PD->getTypeRepr()->getWithoutParens();
+  if (auto *STR = dyn_cast<SpecifierTypeRepr>(typeRepr)) {
+    typeRepr = STR->getBase()->getWithoutParens();
   }
 
-  // First, replace the constraint type with the generic parameter type
-  // placeholder.
-  Diag.fixItReplace(TyReplacementRange, GPNamePlaceholder);
+  SourceRange anyRange;
+  TypeRepr *constraintRepr = typeRepr;
+  if (auto *existentialRepr = dyn_cast<ExistentialTypeRepr>(typeRepr)) {
+    constraintRepr = existentialRepr->getConstraint()->getWithoutParens();
+    auto anyStart = existentialRepr->getAnyLoc();
+    auto anyEnd = existentialRepr->getConstraint()->getStartLoc();
+    anyRange = SourceRange(anyStart, anyEnd);
+  }
 
-  // Remove 'any' if needed, using a character-based removal to pick up
-  // whitespaces between it and its constraint repr.
-  if (RemoveAnyRange.isValid()) {
-    Diag.fixItRemoveChars(RemoveAnyRange.Start, RemoveAnyRange.End);
+  bool needsParens = false;
+  while (auto *metatype = dyn_cast<MetatypeTypeRepr>(constraintRepr)) {
+    // The generic equivalent of 'any P.Type' is '(some P).Type'
+    constraintRepr = metatype->getBase()->getWithoutParens();
+    if (isa<SimpleIdentTypeRepr>(constraintRepr))
+      needsParens = !isa<TupleTypeRepr>(metatype->getBase());
   }
 
-  // Finally, insert the generic parameter declaration.
-  Diag.fixItInsert(GPDeclLoc, GPDeclStr);
+  std::string fix;
+  llvm::raw_string_ostream OS(fix);
+  if (needsParens)
+    OS << "(";
+  OS << "some ";
+  constraintRepr->print(OS);
+  if (needsParens)
+    OS << ")";
+
+  // When removing 'any', use a character-based removal to pick up
+  // whitespaces between it and its constraint repr.
+  Diag
+    .fixItReplace(constraintRepr->getSourceRange(), fix)
+    .fixItRemoveChars(anyRange.Start, anyRange.End);
 
   return true;
 }

test/AutoDiff/validation-test/sr16094.swift

@@ -0,0 +1,125 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import _Differentiation
+import StdlibUnittest
+
+var PullbackTests = TestSuite("Pullback")
+
+extension Dictionary: Differentiable where Value: Differentiable {
+    public typealias TangentVector = [Key: Value.TangentVector]
+    public mutating func move(by direction: TangentVector) {
+        for (componentKey, componentDirection) in direction {
+            func fatalMissingComponent() -> Value {
+                fatalError("missing component \(componentKey) in moved Dictionary")
+            }
+            self[componentKey, default: fatalMissingComponent()].move(by: componentDirection)
+        }
+    }
+    
+    public var zeroTangentVectorInitializer: () -> TangentVector {
+        let listOfKeys = self.keys // capturing only what's needed, not the entire self, in order to not waste memory
+        func initializer() -> Self.TangentVector {
+            return listOfKeys.reduce(into: [Key: Value.TangentVector]()) { $0[$1] = Value.TangentVector.zero }
+        }
+        return initializer
+    }
+}
+
+extension Dictionary: AdditiveArithmetic where Value: AdditiveArithmetic {
+    public static func + (_ lhs: Self, _ rhs: Self) -> Self {
+        return lhs.merging(rhs, uniquingKeysWith: +)
+    }
+
+    public static func - (_ lhs: Self, _ rhs: Self) -> Self {
+        return lhs.merging(rhs.mapValues { .zero - $0 }, uniquingKeysWith: +)
+    }
+
+    public static var zero: Self { [:] }
+}
+
+extension Dictionary where Value: Differentiable {
+    // get
+    @usableFromInline
+    @derivative(of: subscript(_:))
+    func vjpSubscriptGet(key: Key) -> (value: Value?, pullback: (Optional<Value>.TangentVector) -> Dictionary<Key, Value>.TangentVector) {
+        // When adding two dictionaries, nil values are equivalent to zeroes, so there is no need to manually zero-out
+        // every key's value. Instead, it is faster to create a dictionary with the single non-zero entry.
+        return (self[key], { v in
+            if let value = v.value {
+                return [key: value]
+            }
+            else {
+                return .zero
+            }
+        })
+    }
+ }
+
+public extension Dictionary where Value: Differentiable {
+    @differentiable(reverse)
+    mutating func set(_ key: Key, to newValue: Value) {
+        self[key] = newValue
+    }
+
+    @derivative(of: set)
+    mutating func vjpUpdated(_ key: Key, to newValue: Value) -> (value: Void, pullback: (inout TangentVector) -> (Value.TangentVector)) {
+        self.set(key, to: newValue)
+        
+        let forwardCount = self.count
+        let forwardKeys = self.keys // may be heavy to capture all of these, not sure how to do without them though
+        
+        return ((), { v in
+            // manual zero tangent initialization
+            if v.count < forwardCount {
+                v = Self.TangentVector()
+                forwardKeys.forEach { v[$0] = .zero }
+            }
+            
+            if let dElement = v[key] {
+                v[key] = .zero
+                return dElement
+            }
+            else { // should this fail?
+                v[key] = .zero
+                return .zero
+            }
+        })
+    }
+}
+
+
+PullbackTests.test("ConcreteType") {
+  func getD(from newValues: [String: Double], at key: String) -> Double? {
+    if newValues.keys.contains(key) {
+      return newValues[key]
+    }
+    return nil
+  }
+  
+  @differentiable(reverse)
+  func testFunctionD(newValues: [String: Double]) -> Double {
+    return getD(from: newValues, at: "s1")!
+  }
+
+  expectEqual(pullback(at: ["s1": 1.0], of: testFunctionD)(2), ["s1" : 2.0])
+}
+
+PullbackTests.test("GenericType") {
+  func getG<DataType>(from newValues: [String: DataType], at key: String) -> DataType?
+  where DataType: Differentiable {
+    if newValues.keys.contains(key) {
+      return newValues[key]
+    }
+    return nil
+  }
+
+  @differentiable(reverse)
+  func testFunctionG(newValues: [String: Double]) -> Double {
+    return getG(from: newValues, at: "s1")!
+  }
+
+  expectEqual(pullback(at: ["s1": 1.0], of: testFunctionG)(2), ["s1" : 2.0])
+}
+
+runAllTests()