From 75d78304d219cd46ab22e45575d80ab78a1462df Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Mon, 18 Nov 2019 21:21:31 -0800 Subject: [PATCH 1/2] [AutoDiff] Fix variedness propagation for `apply` inout arguments. Propagate variedness from `apply` argument operands to `apply` inout arguments (representing results). `apply` inout arguments are now correctly marked as active, triggering non-differentiability errors. Add `ApplyInstBase::getInoutArguments` for iterating over `@inout` and `@inout_aliasable` arguments. Add non-differentiability diagnostics and activity info tests. Resolves TF-974. --- include/swift/SIL/SILInstruction.h | 31 ++++++++++ .../Mandatory/Differentiation.cpp | 58 ++++++++----------- lib/SILOptimizer/Mandatory/Differentiation.h | 1 - test/AutoDiff/activity_analysis.swift | 49 +++++++++++++++- ...ifferentiation_transform_diagnostics.swift | 55 ++++++++++++++++-- test/AutoDiff/forward_mode_diagnostics.swift | 28 ++++++--- 6 files changed, 173 insertions(+), 49 deletions(-) diff --git a/include/swift/SIL/SILInstruction.h b/include/swift/SIL/SILInstruction.h index 330179ce72a5a..79219b676ff21 100644 --- a/include/swift/SIL/SILInstruction.h +++ b/include/swift/SIL/SILInstruction.h @@ -2046,6 +2046,11 @@ class ApplyInstBase ApplyInstBase(As &&...args) : ApplyInstBase(std::forward(args)...) {} +// SWIFT_ENABLE_TENSORFLOW +private: + const Impl &asImpl() const { return static_cast(*this); } +// SWIFT_ENABLE_TENSORFLOW END + public: using super::getCallee; using super::getSubstCalleeType; @@ -2133,6 +2138,32 @@ class ApplyInstBase bool hasSemantics(StringRef semanticsString) const { return doesApplyCalleeHaveSemantics(getCallee(), semanticsString); } + +// SWIFT_ENABLE_TENSORFLOW +private: + /// Predicate used to filter InoutArgumentRange. + struct OperandToInoutArgument { + const Impl &inst; + OperandToInoutArgument(const Impl &inst) : inst(inst) {} + Optional operator()(unsigned long i) const { + auto paramInfo = inst.getSubstCalleeConv().getParameters()[i]; + auto argument = inst.getArgumentsWithoutIndirectResults()[i]; + if (paramInfo.isIndirectMutating()) + return argument; + return None; + } + }; + +public: + using InoutArgumentRange = + OptionalTransformRange, OperandToInoutArgument>; + /// Returns all `@inout` and `@inout_aliasable` arguments passed to the + /// instruction. + InoutArgumentRange getInoutArguments() const { + return InoutArgumentRange(indices(getArgumentsWithoutIndirectResults()), + OperandToInoutArgument(asImpl())); + } +// SWIFT_ENABLE_TENSORFLOW END }; /// ApplyInst - Represents the full application of a function value. diff --git a/lib/SILOptimizer/Mandatory/Differentiation.cpp b/lib/SILOptimizer/Mandatory/Differentiation.cpp index 0aeb0b9abf66b..55e89f971f5fb 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.cpp +++ b/lib/SILOptimizer/Mandatory/Differentiation.cpp @@ -1622,11 +1622,8 @@ LinearMapInfo::LinearMapInfo(ADContext &context, /// active argument. bool LinearMapInfo::shouldDifferentiateApplyInst(ApplyInst *ai) { // Function applications with an inout argument should be differentiated. - auto paramInfos = ai->getSubstCalleeConv().getParameters(); - auto arguments = ai->getArgumentsWithoutIndirectResults(); - for (auto i : swift::indices(paramInfos)) - if (paramInfos[i].isIndirectInOut() && - activityInfo.isActive(arguments[i], indices)) + for (auto inoutArg : ai->getInoutArguments()) + if (activityInfo.isActive(inoutArg, indices)) return true; bool hasActiveDirectResults = false; @@ -1642,6 +1639,7 @@ bool LinearMapInfo::shouldDifferentiateApplyInst(ApplyInst *ai) { if (isArrayLiteralIntrinsic(ai) && hasActiveResults) return true; + auto arguments = ai->getArgumentsWithoutIndirectResults(); bool hasActiveArguments = llvm::any_of(arguments, [&](SILValue arg) { return activityInfo.isActive(arg, indices); }); return hasActiveResults && hasActiveArguments; @@ -1834,20 +1832,13 @@ void LinearMapInfo::generateDifferentiationDataStructures( for (auto &origBB : *original) { for (auto &inst : origBB) { if (auto *ai = dyn_cast(&inst)) { - // Check for active 'inout' arguments. - bool isInout = false; - auto paramInfos = ai->getSubstCalleeConv().getParameters(); - for (unsigned i : swift::indices(paramInfos)) { - if (paramInfos[i].isIndirectInOut() && - activityInfo.isActive(ai->getArgumentsWithoutIndirectResults()[i], - indices)) { - // Reject functions with active inout arguments. It's not yet - // supported. - isInout = true; - break; - } - } - if (isInout) + // Skip `apply` instructions with active `inout` arguments. + // TODO(TF-129): Support `inout` argument differentiation. + bool hasActiveInoutArgument = + llvm::any_of(ai->getInoutArguments(), [&](SILValue inoutArg) { + return activityInfo.isActive(inoutArg, indices); + }); + if (hasActiveInoutArgument) continue; // Add linear map field to struct for active `apply` instructions. @@ -2008,10 +1999,13 @@ void DifferentiableActivityInfo::propagateVaried( // If callee is non-varying, skip. if (isWithoutDerivative(ai->getCallee())) return; - // If operand is varied, set all direct and indirect results as varied. + // If operand is varied, set all direct/indirect results and inout arguments + // as varied. if (isVaried(operand->get(), i)) { for (auto indRes : ai->getIndirectSILResults()) propagateVariedInwardsThroughProjections(indRes, i); + for (auto inoutArg : ai->getInoutArguments()) + propagateVariedInwardsThroughProjections(inoutArg, i); forEachApplyDirectResult(ai, [&](SILValue directResult) { setVariedAndPropagateToUsers(directResult, i); }); @@ -3778,7 +3772,7 @@ class VJPEmitter final sei->getLoc(), getOpValue(sei->getOperand()), newDefaultBB, caseBBs); } - // If an `apply` has active results or active inout parameters, replace it + // If an `apply` has active results or active inout arguments, replace it // with an `apply` of its VJP. void visitApplyInst(ApplyInst *ai) { // If the function should not be differentiated or its the array literal @@ -3790,13 +3784,10 @@ class VJPEmitter final return; } - // Check and reject functions with active inout arguments. It's not yet - // supported. - auto paramInfos = ai->getSubstCalleeConv().getParameters(); - auto paramArgs = ai->getArgumentsWithoutIndirectResults(); - for (unsigned i : swift::indices(paramInfos)) { - if (paramInfos[i].isIndirectInOut() && - activityInfo.isActive(paramArgs[i], getIndices())) { + // Diagnose functions with active inout arguments. + // TODO(TF-129): Support `inout` argument differentiation. + for (auto inoutArg : ai->getInoutArguments()) { + if (activityInfo.isActive(inoutArg, getIndices())) { context.emitNondifferentiabilityError(ai, invoker, diag::autodiff_cannot_differentiate_through_inout_arguments); errorOccurred = true; @@ -5472,13 +5463,10 @@ class JVPEmitter final return; } - // Check and reject functions with active inout arguments. It's not yet - // supported. - auto paramInfos = ai->getSubstCalleeConv().getParameters(); - auto paramArgs = ai->getArgumentsWithoutIndirectResults(); - for (unsigned i : swift::indices(paramInfos)) { - if (paramInfos[i].isIndirectInOut() && - activityInfo.isActive(paramArgs[i], getIndices())) { + // Diagnose functions with active inout arguments. + // TODO(TF-129): Support `inout` argument differentiation. + for (auto inoutArg : ai->getInoutArguments()) { + if (activityInfo.isActive(inoutArg, getIndices())) { context.emitNondifferentiabilityError(ai, invoker, diag::autodiff_cannot_differentiate_through_inout_arguments); errorOccurred = true; diff --git a/lib/SILOptimizer/Mandatory/Differentiation.h b/lib/SILOptimizer/Mandatory/Differentiation.h index 404338e657ee1..ed1f885e8a50d 100644 --- a/lib/SILOptimizer/Mandatory/Differentiation.h +++ b/lib/SILOptimizer/Mandatory/Differentiation.h @@ -103,7 +103,6 @@ inline void createEntryArguments(SILFunction *f) { auto *decl = new (ctx) ParamDecl(loc, loc, Identifier(), loc, Identifier(), moduleDecl); decl->setSpecifier(ParamDecl::Specifier::Default); -// decl->setType(type.getASTType()); entry->createFunctionArgument(type, decl); }; for (auto indResTy : conv.getIndirectSILResultTypes()) diff --git a/test/AutoDiff/activity_analysis.swift b/test/AutoDiff/activity_analysis.swift index aa153d601244f..50f4744583951 100644 --- a/test/AutoDiff/activity_analysis.swift +++ b/test/AutoDiff/activity_analysis.swift @@ -1,4 +1,4 @@ -// RUN: %target-swift-emit-sil -Xllvm -debug-only=differentiation 2>&1 %s | %FileCheck %s +// RUN: %target-swift-emit-sil -verify -Xllvm -debug-only=differentiation 2>&1 %s | %FileCheck %s // Check that `@noDerivative` struct projections have "NONE" activity. @@ -203,3 +203,50 @@ func TF_954(_ x: Float) -> Float { // CHECK: bb5: // CHECK: [ACTIVE] %40 = begin_access [read] [static] %2 : $*Float // CHECK: [ACTIVE] %41 = load [trivial] %40 : $*Float + +//===----------------------------------------------------------------------===// +// Non-differentiable functions +//===----------------------------------------------------------------------===// + +// Check `inout` arguments. + +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArg(_ x: Float) -> Float { + var result = x + // expected-note @+1 {{cannot differentiate through 'inout' arguments}} + result += x + return result +} + +// CHECK-LABEL: [AD] Activity info for ${{.*}}activeInoutArg{{.*}} at (source=0 parameters=(0)) +// CHECK: [ACTIVE] %0 = argument of bb0 : $Float +// CHECK: [ACTIVE] %2 = alloc_stack $Float, var, name "result" +// CHECK: [ACTIVE] %5 = begin_access [modify] [static] %2 : $*Float +// CHECK: [NONE] // function_ref static Float.+= infix(_:_:) +// CHECK: [NONE] %7 = apply %6(%5, %0, %4) : $@convention(method) (@inout Float, Float, @thin Float.Type) -> () +// CHECK: [ACTIVE] %9 = begin_access [read] [static] %2 : $*Float +// CHECK: [ACTIVE] %10 = load [trivial] %9 : $*Float + +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArgNonactiveInitialResult(_ x: Float) -> Float { + var result: Float = 1 + // expected-note @+1 {{cannot differentiate through 'inout' arguments}} + result += x + return result +} + +// CHECK-LABEL: [AD] Activity info for ${{.*}}activeInoutArgNonactiveInitialResult{{.*}} at (source=0 parameters=(0)) +// CHECK-LABEL: [ACTIVE] %0 = argument of bb0 : $Float +// CHECK-LABEL: [ACTIVE] %2 = alloc_stack $Float, var, name "result" +// CHECK-LABEL: [NONE] // function_ref Float.init(_builtinIntegerLiteral:) +// CHECK-LABEL: [USEFUL] %6 = apply %5(%3, %4) : $@convention(method) (Builtin.IntLiteral, @thin Float.Type) -> Float +// CHECK-LABEL: [USEFUL] %8 = metatype $@thin Float.Type +// CHECK-LABEL: [ACTIVE] %9 = begin_access [modify] [static] %2 : $*Float +// CHECK-LABEL: [NONE] // function_ref static Float.+= infix(_:_:) +// CHECK-LABEL: [NONE] %11 = apply %10(%9, %0, %8) : $@convention(method) (@inout Float, Float, @thin Float.Type) -> () +// CHECK-LABEL: [ACTIVE] %13 = begin_access [read] [static] %2 : $*Float +// CHECK-LABEL: [ACTIVE] %14 = load [trivial] %13 : $*Float diff --git a/test/AutoDiff/differentiation_transform_diagnostics.swift b/test/AutoDiff/differentiation_transform_diagnostics.swift index 60672bf5f743f..2a02804366675 100644 --- a/test/AutoDiff/differentiation_transform_diagnostics.swift +++ b/test/AutoDiff/differentiation_transform_diagnostics.swift @@ -270,23 +270,68 @@ func roundingGivesError(x: Float) -> Float { // Inout arguments //===----------------------------------------------------------------------===// +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} func activeInoutArg(_ x: Float) -> Float { - var a = x + var result = x // expected-note @+1 {{cannot differentiate through 'inout' arguments}} - a += x - return a + result += x + return result } + // expected-error @+1 {{function is not differentiable}} -_ = pullback(at: .zero, in: activeInoutArg(_:)) +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArgNonactiveInitialResult(_ x: Float) -> Float { + var result: Float = 1 + // expected-note @+1 {{cannot differentiate through 'inout' arguments}} + result += x + return result +} +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} func activeInoutArgTuple(_ x: Float) -> Float { var tuple = (x, x) // expected-note @+1 {{cannot differentiate through 'inout' arguments}} tuple.0 *= x return x * tuple.0 } + // expected-error @+1 {{function is not differentiable}} -_ = pullback(at: .zero, in: activeInoutArgTuple(_:)) +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArgControlFlow(_ array: [Float]) -> Float { + var result: Float = 1 + for i in withoutDerivative(at: array).indices { + // expected-note @+1 {{cannot differentiate through 'inout' arguments}} + result += array[i] + } + return result +} + +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArgControlFlowComplex(_ array: [Float], _ bool: Bool) -> Float { + var result: Float = 1 + if bool { + if bool {} + for i in withoutDerivative(at: array).indices { + switch i % 2 { + case 0: continue + case 1: break + default: break + } + result = result + 1 + // expected-note @+1 {{cannot differentiate through 'inout' arguments}} + result += array[i] + } + } + return result +} //===----------------------------------------------------------------------===// // Non-varied results diff --git a/test/AutoDiff/forward_mode_diagnostics.swift b/test/AutoDiff/forward_mode_diagnostics.swift index 5d2f8fba86aa5..18e2ed74c8f9c 100644 --- a/test/AutoDiff/forward_mode_diagnostics.swift +++ b/test/AutoDiff/forward_mode_diagnostics.swift @@ -53,23 +53,37 @@ func calls_diff_of_nested(_ x: Float) -> Float { // Inout arguments //===----------------------------------------------------------------------===// -func activeInoutArg(_ x: Float) -> Float { - var a = x +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} +func activeInoutArgNonactiveInitialResult(_ x: Float) -> Float { + var result: Float = 1 // expected-note @+1 {{cannot differentiate through 'inout' arguments}} - a += x - return a + result += x + return result } -// expected-error @+1 {{function is not differentiable}} -_ = differential(at: .zero, in: activeInoutArg(_:)) +// expected-error @+1 {{function is not differentiable}} +@differentiable +// expected-note @+1 {{when differentiating this function definition}} func activeInoutArgTuple(_ x: Float) -> Float { var tuple = (x, x) // expected-note @+1 {{cannot differentiate through 'inout' arguments}} tuple.0 *= x return x * tuple.0 } + // expected-error @+1 {{function is not differentiable}} -_ = differential(at: .zero, in: activeInoutArgTuple(_:)) +@differentiable +// expected-note @+2 {{when differentiating this function definition}} +// expected-note @+1 {{forward-mode differentiation does not yet support control flow}} +func activeInoutArgControlFlow(_ array: [Float]) -> Float { + var result: Float = 1 + for i in withoutDerivative(at: array).indices { + result += array[i] + } + return result +} //===----------------------------------------------------------------------===// // Non-varied results From d43ec714c6f0bb34b00801b14ba29d69ef63fc43 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Tue, 19 Nov 2019 11:17:11 -0800 Subject: [PATCH 2/2] Store ranges in `OperandToInoutArgument`. --- include/swift/SIL/SILInstruction.h | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/include/swift/SIL/SILInstruction.h b/include/swift/SIL/SILInstruction.h index 79219b676ff21..14b67f0e8af0b 100644 --- a/include/swift/SIL/SILInstruction.h +++ b/include/swift/SIL/SILInstruction.h @@ -2143,13 +2143,16 @@ class ApplyInstBase private: /// Predicate used to filter InoutArgumentRange. struct OperandToInoutArgument { - const Impl &inst; - OperandToInoutArgument(const Impl &inst) : inst(inst) {} + ArrayRef paramInfos; + OperandValueArrayRef arguments; + OperandToInoutArgument(const Impl &inst) + : paramInfos(inst.getSubstCalleeConv().getParameters()), + arguments(inst.getArgumentsWithoutIndirectResults()) { + assert(paramInfos.size() == arguments.size()); + } Optional operator()(unsigned long i) const { - auto paramInfo = inst.getSubstCalleeConv().getParameters()[i]; - auto argument = inst.getArgumentsWithoutIndirectResults()[i]; - if (paramInfo.isIndirectMutating()) - return argument; + if (paramInfos[i].isIndirectMutating()) + return arguments[i]; return None; } };