[AutoDiff] Support differentiation of branching cast instructions.

dan-zheng · dan-zheng · commit d5d076db6a23 · 2020-05-28T12:54:38.000-07:00
Support differentiation of `is` and `as?` operators.

These operators lower to branching cast SIL instructions, requiring control
flow differentiation support.

Resolves SR-12898.
diff --git a/include/swift/SILOptimizer/Differentiation/VJPEmitter.h b/include/swift/SILOptimizer/Differentiation/VJPEmitter.h
@@ -153,6 +153,12 @@ class VJPEmitter final
 
   void visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai);
 
+  void visitCheckedCastBranchInst(CheckedCastBranchInst *ccbi);
+
+  void visitCheckedCastValueBranchInst(CheckedCastValueBranchInst *ccvbi);
+
+  void visitCheckedCastAddrBranchInst(CheckedCastAddrBranchInst *ccabi);
+
   // If an `apply` has active results or active inout arguments, replace it
   // with an `apply` of its VJP.
   void visitApplyInst(ApplyInst *ai);
diff --git a/lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp b/lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp
@@ -193,12 +193,26 @@ void DifferentiableActivityInfo::propagateVaried(
       if (auto *destBBArg = cbi->getArgForOperand(operand))
         setVariedAndPropagateToUsers(destBBArg, i);
   }
-  // Handle `switch_enum`.
-  else if (auto *sei = dyn_cast<SwitchEnumInst>(inst)) {
-    if (isVaried(sei->getOperand(), i))
-      for (auto *succBB : sei->getSuccessorBlocks())
+  // Handle `checked_cast_addr_br`.
+  // Propagate variedness from source operand to destination operand, in
+  // addition to all successor block arguments.
+  else if (auto *ccabi = dyn_cast<CheckedCastAddrBranchInst>(inst)) {
+    if (isVaried(ccabi->getSrc(), i)) {
+      setVariedAndPropagateToUsers(ccabi->getDest(), i);
+      for (auto *succBB : ccabi->getSuccessorBlocks())
         for (auto *arg : succBB->getArguments())
           setVariedAndPropagateToUsers(arg, i);
+    }
+  }
+  // Handle all other terminators: if any operand is active, propagate
+  // variedness to all successor block arguments. This logic may be incorrect
+  // for some terminator instructions, so special cases must be defined above.
+  else if (auto *termInst = dyn_cast<TermInst>(inst)) {
+    for (auto &op : termInst->getAllOperands())
+      if (isVaried(op.get(), i))
+        for (auto *succBB : termInst->getSuccessorBlocks())
+          for (auto *arg : succBB->getArguments())
+            setVariedAndPropagateToUsers(arg, i);
   }
   // Handle everything else.
   else {
diff --git a/lib/SILOptimizer/Differentiation/VJPEmitter.cpp b/lib/SILOptimizer/Differentiation/VJPEmitter.cpp
@@ -481,6 +481,47 @@ void VJPEmitter::visitSwitchEnumAddrInst(SwitchEnumAddrInst *seai) {
   visitSwitchEnumInstBase(seai);
 }
 
+void VJPEmitter::visitCheckedCastBranchInst(CheckedCastBranchInst *ccbi) {
+  // Build pullback struct value for original block.
+  auto *pbStructVal = buildPullbackValueStructValue(ccbi->getParent());
+  // Create a new `checked_cast_branch` instruction.
+  getBuilder().createCheckedCastBranch(
+      ccbi->getLoc(), ccbi->isExact(), getOpValue(ccbi->getOperand()),
+      getOpType(ccbi->getTargetLoweredType()),
+      getOpASTType(ccbi->getTargetFormalType()),
+      createTrampolineBasicBlock(ccbi, pbStructVal, ccbi->getSuccessBB()),
+      createTrampolineBasicBlock(ccbi, pbStructVal, ccbi->getFailureBB()),
+      ccbi->getTrueBBCount(), ccbi->getFalseBBCount());
+}
+
+void VJPEmitter::visitCheckedCastValueBranchInst(
+    CheckedCastValueBranchInst *ccvbi) {
+  // Build pullback struct value for original block.
+  auto *pbStructVal = buildPullbackValueStructValue(ccvbi->getParent());
+  // Create a new `checked_cast_value_branch` instruction.
+  getBuilder().createCheckedCastValueBranch(
+      ccvbi->getLoc(), getOpValue(ccvbi->getOperand()),
+      getOpASTType(ccvbi->getSourceFormalType()),
+      getOpType(ccvbi->getTargetLoweredType()),
+      getOpASTType(ccvbi->getTargetFormalType()),
+      createTrampolineBasicBlock(ccvbi, pbStructVal, ccvbi->getSuccessBB()),
+      createTrampolineBasicBlock(ccvbi, pbStructVal, ccvbi->getFailureBB()));
+}
+
+void VJPEmitter::visitCheckedCastAddrBranchInst(
+    CheckedCastAddrBranchInst *ccabi) {
+  // Build pullback struct value for original block.
+  auto *pbStructVal = buildPullbackValueStructValue(ccabi->getParent());
+  // Create a new `checked_cast_addr_branch` instruction.
+  getBuilder().createCheckedCastAddrBranch(
+      ccabi->getLoc(), ccabi->getConsumptionKind(), getOpValue(ccabi->getSrc()),
+      getOpASTType(ccabi->getSourceFormalType()), getOpValue(ccabi->getDest()),
+      getOpASTType(ccabi->getTargetFormalType()),
+      createTrampolineBasicBlock(ccabi, pbStructVal, ccabi->getSuccessBB()),
+      createTrampolineBasicBlock(ccabi, pbStructVal, ccabi->getFailureBB()),
+      ccabi->getTrueBBCount(), ccabi->getFalseBBCount());
+}
+
 void VJPEmitter::visitApplyInst(ApplyInst *ai) {
   // If callee should not be differentiated, do standard cloning.
   if (!pullbackInfo.shouldDifferentiateApplySite(ai)) {
diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp
@@ -152,10 +152,12 @@ static bool diagnoseUnsupportedControlFlow(ADContext &context,
   // Diagnose unsupported branching terminators.
   for (auto &bb : *original) {
     auto *term = bb.getTerminator();
-    // Supported terminators are: `br`, `cond_br`, `switch_enum`,
-    // `switch_enum_addr`.
+    // Check supported branching terminators.
     if (isa<BranchInst>(term) || isa<CondBranchInst>(term) ||
-        isa<SwitchEnumInst>(term) || isa<SwitchEnumAddrInst>(term))
+        isa<SwitchEnumInst>(term) || isa<SwitchEnumAddrInst>(term) ||
+        isa<CheckedCastBranchInst>(term) ||
+        isa<CheckedCastValueBranchInst>(term) ||
+        isa<CheckedCastAddrBranchInst>(term))
       continue;
     // If terminator is an unsupported branching terminator, emit an error.
     if (term->isBranch()) {
diff --git a/test/AutoDiff/SILOptimizer/activity_analysis.swift b/test/AutoDiff/SILOptimizer/activity_analysis.swift
@@ -122,6 +122,116 @@ func TF_954(_ x: Float) -> Float {
 // CHECK: [ACTIVE]   %40 = begin_access [read] [static] %2 : $*Float
 // CHECK: [ACTIVE]   %41 = load [trivial] %40 : $*Float
 
+//===----------------------------------------------------------------------===//
+// Branching cast instructions
+//===----------------------------------------------------------------------===//
+
+@differentiable
+func checked_cast_branch(_ x: Float) -> Float {
+  // expected-warning @+1 {{'is' test is always true}}
+  if Int.self is Any.Type {
+    return x + x
+  }
+  return x * x
+}
+
+// CHECK-LABEL: [AD] Activity info for ${{.*}}checked_cast_branch{{.*}} at (source=0 parameters=(0))
+// CHECK: bb0:
+// CHECK: [ACTIVE] %0 = argument of bb0 : $Float
+// CHECK: [NONE]   %2 = metatype $@thin Int.Type
+// CHECK: [NONE]   %3 = metatype $@thick Int.Type
+// CHECK: bb1:
+// CHECK: [NONE] %5 = argument of bb1 : $@thick Any.Type
+// CHECK: [NONE]   %6 = integer_literal $Builtin.Int1, -1
+// CHECK: bb2:
+// CHECK: [NONE] %8 = argument of bb2 : $@thick Int.Type
+// CHECK: [NONE]   %9 = integer_literal $Builtin.Int1, 0
+// CHECK: bb3:
+// CHECK: [NONE] %11 = argument of bb3 : $Builtin.Int1
+// CHECK: [NONE]   %12 = metatype $@thin Bool.Type
+// CHECK: [NONE]   // function_ref Bool.init(_builtinBooleanLiteral:)
+// CHECK: [NONE]   %14 = apply %13(%11, %12) : $@convention(method) (Builtin.Int1, @thin Bool.Type) -> Bool
+// CHECK: [NONE]   %15 = struct_extract %14 : $Bool, #Bool._value
+// CHECK: bb4:
+// CHECK: [USEFUL]   %17 = metatype $@thin Float.Type
+// CHECK: [NONE]   // function_ref static Float.+ infix(_:_:)
+// CHECK: [ACTIVE]   %19 = apply %18(%0, %0, %17) : $@convention(method) (Float, Float, @thin Float.Type) -> Float
+// CHECK: bb5:
+// CHECK: [USEFUL]   %21 = metatype $@thin Float.Type
+// CHECK: [NONE]   // function_ref static Float.* infix(_:_:)
+// CHECK: [ACTIVE]   %23 = apply %22(%0, %0, %21) : $@convention(method) (Float, Float, @thin Float.Type) -> Float
+
+// CHECK-LABEL: sil hidden [ossa] @${{.*}}checked_cast_branch{{.*}} : $@convention(thin) (Float) -> Float {
+// CHECK:   checked_cast_br %3 : $@thick Int.Type to Any.Type, bb1, bb2
+// CHECK: }
+
+@differentiable
+func checked_cast_addr_nonactive_result<T: Differentiable>(_ x: T) -> T {
+  if let _ = x as? Float {
+    // Do nothing with `y: Float?` value.
+  }
+  return x
+}
+
+// CHECK-LABEL: [AD] Activity info for ${{.*}}checked_cast_addr_nonactive_result{{.*}} at (source=0 parameters=(0))
+// CHECK: bb0:
+// CHECK: [ACTIVE] %0 = argument of bb0 : $*T
+// CHECK: [ACTIVE] %1 = argument of bb0 : $*T
+// CHECK: [VARIED]   %3 = alloc_stack $T
+// CHECK: [VARIED]   %5 = alloc_stack $Float
+// CHECK: bb1:
+// CHECK: [VARIED]   %7 = load [trivial] %5 : $*Float
+// CHECK: [VARIED]   %8 = enum $Optional<Float>, #Optional.some!enumelt, %7 : $Float
+// CHECK: bb2:
+// CHECK: [NONE]   %11 = enum $Optional<Float>, #Optional.none!enumelt
+// CHECK: bb3:
+// CHECK: [VARIED] %14 = argument of bb3 : $Optional<Float>
+// CHECK: bb4:
+// CHECK: bb5:
+// CHECK: [VARIED] %18 = argument of bb5 : $Float
+// CHECK: bb6:
+// CHECK: [NONE]   %22 = tuple ()
+
+// CHECK-LABEL: sil hidden [ossa] @${{.*}}checked_cast_addr_nonactive_result{{.*}} : $@convention(thin) <T where T : Differentiable> (@in_guaranteed T) -> @out T {
+// CHECK:   checked_cast_addr_br take_always T in %3 : $*T to Float in %5 : $*Float, bb1, bb2
+// CHECK: }
+
+// expected-error @+1 {{function is not differentiable}}
+@differentiable
+// expected-note @+1 {{when differentiating this function definition}}
+func checked_cast_addr_active_result<T: Differentiable>(x: T) -> T {
+  // expected-note @+1 {{differentiating enum values is not yet supported}}
+  if let y = x as? Float {
+    // Use `y: Float?` value in an active way.
+    return y as! T
+  }
+  return x
+}
+
+// CHECK-LABEL: [AD] Activity info for ${{.*}}checked_cast_addr_active_result{{.*}} at (source=0 parameters=(0))
+// CHECK: bb0:
+// CHECK: [ACTIVE] %0 = argument of bb0 : $*T
+// CHECK: [ACTIVE] %1 = argument of bb0 : $*T
+// CHECK: [ACTIVE]   %3 = alloc_stack $T
+// CHECK: [ACTIVE]   %5 = alloc_stack $Float
+// CHECK: bb1:
+// CHECK: [ACTIVE]   %7 = load [trivial] %5 : $*Float
+// CHECK: [ACTIVE]   %8 = enum $Optional<Float>, #Optional.some!enumelt, %7 : $Float
+// CHECK: bb2:
+// CHECK: [USEFUL]   %11 = enum $Optional<Float>, #Optional.none!enumelt
+// CHECK: bb3:
+// CHECK: [ACTIVE] %14 = argument of bb3 : $Optional<Float>
+// CHECK: bb4:
+// CHECK: [ACTIVE] %16 = argument of bb4 : $Float
+// CHECK: [ACTIVE]   %19 = alloc_stack $Float
+// CHECK: bb5:
+// CHECK: bb6:
+// CHECK: [NONE]   %27 = tuple ()
+
+// CHECK-LABEL: sil hidden [ossa] @${{.*}}checked_cast_addr_active_result{{.*}} : $@convention(thin) <T where T : Differentiable> (@in_guaranteed T) -> @out T {
+// CHECK:   checked_cast_addr_br take_always T in %3 : $*T to Float in %5 : $*Float, bb1, bb2
+// CHECK: }
+
 //===----------------------------------------------------------------------===//
 // Array literal differentiation
 //===----------------------------------------------------------------------===//
diff --git a/test/AutoDiff/validation-test/control_flow.swift b/test/AutoDiff/validation-test/control_flow.swift
@@ -715,4 +715,25 @@ ControlFlowTests.test("Loops") {
   expectEqual((24, 28), valueWithGradient(at: 2, in: { x in nested_loop2(x, count: 4) }))
 }
 
+ControlFlowTests.test("BranchingCastInstructions") {
+  // checked_cast_br
+  func typeCheckOperator<T>(_ x: Float, _ metatype: T.Type) -> Float {
+    if metatype is Int.Type {
+      return x + x
+    }
+    return x * x
+  }
+  expectEqual((6, 2), valueWithGradient(at: 3, in: { typeCheckOperator($0, Int.self) }))
+  expectEqual((9, 6), valueWithGradient(at: 3, in: { typeCheckOperator($0, Float.self) }))
+
+  // checked_cast_addr_br
+  func conditionalCast<T: Differentiable>(_ x: T) -> T {
+    if let _ = x as? Float {
+      // Do nothing with `y: Float?` value.
+    }
+    return x
+  }
+  expectEqual((3, 1), valueWithGradient(at: Float(3), in: conditionalCast))
+}
+
 runAllTests()
