From 65e1ae477c07ad962f257333e9eb7e0e75e79fdc Mon Sep 17 00:00:00 2001 From: Marc Rasi Date: Thu, 6 Dec 2018 15:57:51 -0800 Subject: [PATCH 1/2] fix VJP differentiation for TensorFlow autodiff tests --- .../Mandatory/TFDifferentiation.cpp | 17 +++++++++++++++++ lib/SILOptimizer/Mandatory/TFLowerGraph.cpp | 2 ++ stdlib/public/TensorFlow/Ops.swift | 5 ++++- test/TensorFlow/tensor_autodiff.swift | 1 + .../tensor_autodiff_runtime.swift | 1 + 5 files changed, 25 insertions(+), 1 deletion(-) diff --git a/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp b/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp index 6253491e2023e..5be10e5cbb672 100644 --- a/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp +++ b/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp @@ -4466,6 +4466,23 @@ void DifferentiationTask::createVJP() { SILDebugScope(original->getLocation(), vjp)); attr->setVJPName(vjp->getName()); + // Work around a bad interaction between VJPs, TFDeabstraction, and SIL + // optimizations. + // + // The bad interaction is: TFDeabstraction cannot inline the functions that + // the VJP partially applies. When run in "-Onone" mode, this is fine (but + // inefficient), because we just generate sends/recvs to handle it. But in + // "-O" mode, later SIL optimization passes: + // - fold the partial applications and inline them, or + // - specialize the partial_apply callees. + // In either case, the inlined/specialized bodies haven't been deabstracted, + // so partitioning crashes when it sees them. + // + // TODO: Fix this problem in a way that allows inlining/specialization of + // VJPs. Teaching TFDeabstraction to fold partial applications might work. + vjp->setInlineStrategy(NoInline); + vjp->addSemanticsAttr("optimize.sil.specialize.generic.never"); + LLVM_DEBUG(llvm::dbgs() << " vjp type: " << vjp->getLoweredFunctionType() << "\n"); diff --git a/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp b/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp index e64abeb111b3e..6c3b6be61a21c 100644 --- a/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp +++ b/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp @@ -2694,6 +2694,8 @@ bool TFGraphLowering::serializeGraphProtoBuf(ASTContext &ctx, /// repl) std::string getTFCompatibleFuncName(SILFunction *fn) { auto fnName = fn->getName(); + if (fnName.startswith("AD__")) + fnName = fnName.substr(4); if (fnName.startswith("$")) fnName = fnName.substr(1); diff --git a/stdlib/public/TensorFlow/Ops.swift b/stdlib/public/TensorFlow/Ops.swift index 097c487472bab..75ffc84617428 100644 --- a/stdlib/public/TensorFlow/Ops.swift +++ b/stdlib/public/TensorFlow/Ops.swift @@ -275,7 +275,10 @@ public extension Tensor where Scalar : Numeric { public func matmul( _ left: Tensor, _ right: Tensor ) -> Tensor { - return Raw.matMul(left, right) + // Default arguments specified explicitly to avoid "external declarations of + // SILFunctions with shared visibility is not allowed" SILVerifier error in + // "tests/AutoDiff/tensor_autodiff_runtime.swift". + return Raw.matMul(left, right, transposeA: false, transposeB: false) } infix operator • : MultiplicationPrecedence diff --git a/test/TensorFlow/tensor_autodiff.swift b/test/TensorFlow/tensor_autodiff.swift index 9944117f0d42d..419f6b1b67a57 100644 --- a/test/TensorFlow/tensor_autodiff.swift +++ b/test/TensorFlow/tensor_autodiff.swift @@ -1,4 +1,5 @@ // RUN: %target-swift-frontend -O -emit-sil %s | %FileCheck %s +// RUN: %target-swift-frontend -Xllvm -differentiation-use-vjp -O -emit-sil %s | %FileCheck %s import TensorFlow diff --git a/test/TensorFlowRuntime/tensor_autodiff_runtime.swift b/test/TensorFlowRuntime/tensor_autodiff_runtime.swift index 7d274987e04f2..3b53e3fee9a7c 100644 --- a/test/TensorFlowRuntime/tensor_autodiff_runtime.swift +++ b/test/TensorFlowRuntime/tensor_autodiff_runtime.swift @@ -1,4 +1,5 @@ // RUN: %target-run-simple-swift %swift-tensorflow-test-run-extra-options +// RUN: %target-run-use-vjp-swift %swift-tensorflow-test-run-extra-options // // TODO(SR-9110): Make this pass in dynamic compilation mode. // %target-run-dynamic-compilation-swift From d10bdd9547d57a6d105473d6b18102f034f0b4fd Mon Sep 17 00:00:00 2001 From: Marc Rasi Date: Tue, 11 Dec 2018 21:26:30 -0800 Subject: [PATCH 2/2] fix test --- lib/SILOptimizer/Mandatory/TFDifferentiation.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp b/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp index 5be10e5cbb672..b1e84686c3641 100644 --- a/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp +++ b/lib/SILOptimizer/Mandatory/TFDifferentiation.cpp @@ -4489,6 +4489,7 @@ void DifferentiationTask::createVJP() { // We'll use these conventions frequently. auto originalConv = original->getConventions(); auto primalConv = primal->getConventions(); + auto vjpConv = vjp->getConventions(); // Keep track of some stack allocation to clean up. SmallVector stackAllocsToCleanUp; @@ -4496,7 +4497,6 @@ void DifferentiationTask::createVJP() { // Validate signatures. #ifndef NDEBUG auto adjointConv = adjoint->getConventions(); - auto vjpConv = vjp->getConventions(); unsigned numOriginalParameters = originalConv.getNumParameters(); unsigned numOriginalResults = originalConv.getResults().size(); @@ -4555,7 +4555,16 @@ void DifferentiationTask::createVJP() { // Create the entry block with indirect results and parameters. auto *entry = vjp->createBasicBlock(); - createEntryArguments(vjp); + unsigned argumentIndex = 0; + for (auto indResultTy : vjpConv.getIndirectSILResultTypes()) + entry->createFunctionArgument( + vjp->mapTypeIntoContext(indResultTy).getAddressType(), + original->getEntryBlock()->getArgument(argumentIndex++)->getDecl()); + for (auto paramTy : vjpConv.getParameterSILTypes()) + entry->createFunctionArgument( + vjp->mapTypeIntoContext(paramTy), + original->getEntryBlock()->getArgument(argumentIndex++)->getDecl()); + SILBuilder builder(entry); auto loc = vjp->getLocation();