Merge pull request swiftlang#32031 from dan-zheng/fix-alloc-stack-cloning

dan-zheng · web-flow · commit 8c70a9625879 · 2020-05-29T14:50:08.000-07:00
[AutoDiff] Fix differentiation crashes related to definite initialization.
diff --git a/include/swift/SIL/SILCloner.h b/include/swift/SIL/SILCloner.h
@@ -802,9 +802,9 @@ SILCloner<ImplClass>::visitAllocStackInst(AllocStackInst *Inst) {
     Loc = MandatoryInlinedLocation::getAutoGeneratedLocation();
     VarInfo = None;
   }
-  recordClonedInstruction(Inst,
-                          getBuilder().createAllocStack(
-                              Loc, getOpType(Inst->getElementType()), VarInfo));
+  recordClonedInstruction(Inst, getBuilder().createAllocStack(
+                                    Loc, getOpType(Inst->getElementType()),
+                                    VarInfo, Inst->hasDynamicLifetime()));
 }
 
 template<typename ImplClass>
diff --git a/include/swift/SILOptimizer/Differentiation/VJPEmitter.h b/include/swift/SILOptimizer/Differentiation/VJPEmitter.h
@@ -130,8 +130,9 @@ class VJPEmitter final
                                             StructInst *pbStructVal,
                                             SILBasicBlock *succBB);
 
-  /// Build a pullback struct value for the given original block.
-  StructInst *buildPullbackValueStructValue(SILBasicBlock *bb);
+  /// Build a pullback struct value for the given original terminator
+  /// instruction.
+  StructInst *buildPullbackValueStructValue(TermInst *termInst);
 
   /// Build a predecessor enum instance using the given builder for the given
   /// original predecessor/successor blocks and pullback struct value.
diff --git a/lib/SILOptimizer/Differentiation/VJPEmitter.cpp b/lib/SILOptimizer/Differentiation/VJPEmitter.cpp
@@ -315,9 +315,10 @@ SILType VJPEmitter::getNominalDeclLoweredType(NominalTypeDecl *nominal) {
   return getLoweredType(nominalType);
 }
 
-StructInst *VJPEmitter::buildPullbackValueStructValue(SILBasicBlock *origBB) {
-  assert(origBB->getParent() == original);
-  auto loc = origBB->getParent()->getLocation();
+StructInst *VJPEmitter::buildPullbackValueStructValue(TermInst *termInst) {
+  assert(termInst->getFunction() == original);
+  auto loc = RegularLocation::getAutoGeneratedLocation();
+  auto origBB = termInst->getParent();
   auto *vjpBB = BBMap[origBB];
   auto *pbStruct = pullbackInfo.getLinearMapStruct(origBB);
   auto structLoweredTy = getNominalDeclLoweredType(pbStruct);
@@ -326,14 +327,15 @@ StructInst *VJPEmitter::buildPullbackValueStructValue(SILBasicBlock *origBB) {
     auto *predEnumArg = vjpBB->getArguments().back();
     bbPullbackValues.insert(bbPullbackValues.begin(), predEnumArg);
   }
+  getBuilder().setCurrentDebugScope(getOpScope(termInst->getDebugScope()));
   return getBuilder().createStruct(loc, structLoweredTy, bbPullbackValues);
 }
 
 EnumInst *VJPEmitter::buildPredecessorEnumValue(SILBuilder &builder,
                                                 SILBasicBlock *predBB,
                                                 SILBasicBlock *succBB,
                                                 SILValue pbStructVal) {
-  auto loc = pbStructVal.getLoc();
+  auto loc = RegularLocation::getAutoGeneratedLocation();
   auto *succEnum = pullbackInfo.getBranchingTraceDecl(succBB);
   auto enumLoweredTy = getNominalDeclLoweredType(succEnum);
   auto *enumEltDecl =
@@ -361,7 +363,7 @@ void VJPEmitter::visitReturnInst(ReturnInst *ri) {
 
   // Build pullback struct value for original block.
   auto *origExit = ri->getParent();
-  auto *pbStructVal = buildPullbackValueStructValue(origExit);
+  auto *pbStructVal = buildPullbackValueStructValue(ri);
 
   // Get the value in the VJP corresponding to the original result.
   auto *origRetInst = cast<ReturnInst>(origExit->getTerminator());
@@ -416,7 +418,7 @@ void VJPEmitter::visitBranchInst(BranchInst *bi) {
   // Build pullback struct value for original block.
   // Build predecessor enum value for destination block.
   auto *origBB = bi->getParent();
-  auto *pbStructVal = buildPullbackValueStructValue(origBB);
+  auto *pbStructVal = buildPullbackValueStructValue(bi);
   auto *enumVal = buildPredecessorEnumValue(getBuilder(), origBB,
                                             bi->getDestBB(), pbStructVal);
 
@@ -433,7 +435,7 @@ void VJPEmitter::visitBranchInst(BranchInst *bi) {
 
 void VJPEmitter::visitCondBranchInst(CondBranchInst *cbi) {
   // Build pullback struct value for original block.
-  auto *pbStructVal = buildPullbackValueStructValue(cbi->getParent());
+  auto *pbStructVal = buildPullbackValueStructValue(cbi);
   // Create a new `cond_br` instruction.
   getBuilder().createCondBranch(
       cbi->getLoc(), getOpValue(cbi->getCondition()),
@@ -443,7 +445,7 @@ void VJPEmitter::visitCondBranchInst(CondBranchInst *cbi) {
 
 void VJPEmitter::visitSwitchEnumInstBase(SwitchEnumInstBase *sei) {
   // Build pullback struct value for original block.
-  auto *pbStructVal = buildPullbackValueStructValue(sei->getParent());
+  auto *pbStructVal = buildPullbackValueStructValue(sei);
 
   // Create trampoline successor basic blocks.
   SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 4> caseBBs;
@@ -483,7 +485,7 @@ void VJPEmitter::visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai) {
 
 void VJPEmitter::visitCheckedCastBranchInst(CheckedCastBranchInst *ccbi) {
   // Build pullback struct value for original block.
-  auto *pbStructVal = buildPullbackValueStructValue(ccbi->getParent());
+  auto *pbStructVal = buildPullbackValueStructValue(ccbi);
   // Create a new `checked_cast_branch` instruction.
   getBuilder().createCheckedCastBranch(
       ccbi->getLoc(), ccbi->isExact(), getOpValue(ccbi->getOperand()),
@@ -497,7 +499,7 @@ void VJPEmitter::visitCheckedCastBranchInst(CheckedCastBranchInst *ccbi) {
 void VJPEmitter::visitCheckedCastValueBranchInst(
     CheckedCastValueBranchInst *ccvbi) {
   // Build pullback struct value for original block.
-  auto *pbStructVal = buildPullbackValueStructValue(ccvbi->getParent());
+  auto *pbStructVal = buildPullbackValueStructValue(ccvbi);
   // Create a new `checked_cast_value_branch` instruction.
   getBuilder().createCheckedCastValueBranch(
       ccvbi->getLoc(), getOpValue(ccvbi->getOperand()),
@@ -511,7 +513,7 @@ void VJPEmitter::visitCheckedCastValueBranchInst(
 void VJPEmitter::visitCheckedCastAddrBranchInst(
     CheckedCastAddrBranchInst *ccabi) {
   // Build pullback struct value for original block.
-  auto *pbStructVal = buildPullbackValueStructValue(ccabi->getParent());
+  auto *pbStructVal = buildPullbackValueStructValue(ccabi);
   // Create a new `checked_cast_addr_branch` instruction.
   getBuilder().createCheckedCastAddrBranch(
       ccabi->getLoc(), ccabi->getConsumptionKind(), getOpValue(ccabi->getSrc()),
diff --git a/test/AutoDiff/compiler_crashers_fixed/sr12886-sr112887-vjp-emitter-definite-initialization.swift b/test/AutoDiff/compiler_crashers_fixed/sr12886-sr112887-vjp-emitter-definite-initialization.swift
@@ -0,0 +1,51 @@
+// RUN: %target-build-swift %s
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+// Test crashes related to differentiation and definite intiialization.
+
+// SR-12886: SIL memory lifetime verification error due to
+// `SILCloner::visitAllocStack` not copying the `[dynamic_lifetime]` attribute.
+
+// SR-12887: Debug scope error for pullback struct `struct` instruction
+// generated by `VJPEmitter`.
+
+import _Differentiation
+
+enum Enum {
+  case a
+}
+
+struct Tensor<T>: Differentiable {
+  @noDerivative var x: T
+  @noDerivative var optional: Int?
+
+  init(_ x: T, _ e: Enum) {
+    self.x = x
+    switch e {
+      case .a: optional = 1
+    }
+  }
+
+  // Definite initialization triggers for this initializer.
+  @differentiable
+  init(_ x: T, _ other: Self) {
+    self = Self(x, Enum.a)
+  }
+}
+
+// Check that `allock_stack [dynamic_lifetime]` attribute is correctly cloned.
+
+// CHECK-LABEL: sil hidden @$s4main6TensorVyACyxGx_ADtcfC : $@convention(method) <T> (@in T, @in Tensor<T>, @thin Tensor<T>.Type) -> @out Tensor<T> {
+// CHECK: [[SELF_ALLOC:%.*]] = alloc_stack [dynamic_lifetime] $Tensor<T>, var, name "self"
+
+// CHECK-LABEL: sil hidden @AD__$s4main6TensorVyACyxGx_ADtcfC__vjp_src_0_wrt_1_l : $@convention(method) <τ_0_0> (@in τ_0_0, @in Tensor<τ_0_0>, @thin Tensor<τ_0_0>.Type) -> (@out Tensor<τ_0_0>, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <Tensor<τ_0_0>.TangentVector, Tensor<τ_0_0>.TangentVector>) {
+// CHECK: [[SELF_ALLOC:%.*]] = alloc_stack [dynamic_lifetime] $Tensor<τ_0_0>, var, name "self"
+
+// SR-12886 original error:
+// SIL memory lifetime failure in @AD__$s5crash6TensorVyACyxGx_ADtcfC__vjp_src_0_wrt_1_l: memory is not initialized, but should
+// memory location:   %29 = struct_element_addr %5 : $*Tensor<τ_0_0>, #Tensor.x // user: %30
+// at instruction:   destroy_addr %29 : $*τ_0_0                     // id: %30
+
+// SR-12887 original error:
+// SIL verification failed: Basic block contains a non-contiguous lexical scope at -Onone: DS == LastSeenScope
+// %26 = struct $_AD__$s5crash6TensorVyACyxGx_ADtcfC_bb0__PB__src_0_wrt_1_l<τ_0_0> () // users: %34, %28

Original file line number	Diff line number	Diff line change
`@@ -802,9 +802,9 @@ SILCloner<ImplClass>::visitAllocStackInst(AllocStackInst *Inst) {`
`802`	`802`	`Loc = MandatoryInlinedLocation::getAutoGeneratedLocation();`
`803`	`803`	`VarInfo = None;`
`804`	`804`	`}`
`805`		`- recordClonedInstruction(Inst,`
`806`		`- getBuilder().createAllocStack(`
`807`		`- Loc, getOpType(Inst->getElementType()), VarInfo));`
	`805`	`+ recordClonedInstruction(Inst, getBuilder().createAllocStack(`
	`806`	`+ Loc, getOpType(Inst->getElementType()),`
	`807`	`+ VarInfo, Inst->hasDynamicLifetime()));`
`808`	`808`	`}`
`809`	`809`
`810`	`810`	`template<typename ImplClass>`