From 1df02773f6e4c5457e54fcac9eab9aaf3c830646 Mon Sep 17 00:00:00 2001 From: Ewa Matejska <15254638+ematejska@users.noreply.github.com> Date: Fri, 10 Apr 2020 10:31:23 -0700 Subject: [PATCH 1/3] Autodiff Upstream : Adding DifferentiabilityWitnessDevirtualizer --- .../swift/SILOptimizer/PassManager/Passes.def | 3 + lib/SILOptimizer/PassManager/PassPipeline.cpp | 5 ++ lib/SILOptimizer/Transforms/CMakeLists.txt | 1 + .../DifferentiabilityWitnessDevirtualizer.cpp | 72 +++++++++++++++++++ .../differentiability_witness_inlining.sil | 42 +++++++++++ 5 files changed, 123 insertions(+) create mode 100644 lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp create mode 100644 test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil diff --git a/include/swift/SILOptimizer/PassManager/Passes.def b/include/swift/SILOptimizer/PassManager/Passes.def index 83db39ea944b8..55f12b2d7d8d5 100644 --- a/include/swift/SILOptimizer/PassManager/Passes.def +++ b/include/swift/SILOptimizer/PassManager/Passes.def @@ -152,6 +152,9 @@ PASS(DiagnoseUnreachable, "diagnose-unreachable", "Diagnose Unreachable Code") PASS(DiagnosticConstantPropagation, "diagnostic-constant-propagation", "Constants Propagation for Diagnostics") +PASS(DifferentiabilityWitnessDevirtualizer, + "differentiability-witness-devirtualizer", + "Inlines Differentiability Witnesses") PASS(EagerSpecializer, "eager-specializer", "Eager Specialization via @_specialize") PASS(EarlyCodeMotion, "early-codemotion", diff --git a/lib/SILOptimizer/PassManager/PassPipeline.cpp b/lib/SILOptimizer/PassManager/PassPipeline.cpp index f2d717f4f4f72..926d36aebc56b 100644 --- a/lib/SILOptimizer/PassManager/PassPipeline.cpp +++ b/lib/SILOptimizer/PassManager/PassPipeline.cpp @@ -408,6 +408,11 @@ static void addPerfEarlyModulePassPipeline(SILPassPipelinePlan &P) { // Cleanup after SILGen: remove unneeded borrows/copies. P.addSemanticARCOpts(); + // Devirtualizes differentiability witnesses into functions that reference them. + // This unblocks many other passes' optimizations (e.g. inlining) and this is + // not blocked by any other passes' optimizations, so do it early. + P.addDifferentiabilityWitnessDevirtualizer(); + // Strip ownership from non-transparent functions. if (P.getOptions().StripOwnershipAfterSerialization) P.addNonTransparentFunctionOwnershipModelEliminator(); diff --git a/lib/SILOptimizer/Transforms/CMakeLists.txt b/lib/SILOptimizer/Transforms/CMakeLists.txt index 484d4280bf447..e060c0208c1e6 100644 --- a/lib/SILOptimizer/Transforms/CMakeLists.txt +++ b/lib/SILOptimizer/Transforms/CMakeLists.txt @@ -17,6 +17,7 @@ silopt_register_sources( DeadStoreElimination.cpp DestroyHoisting.cpp Devirtualizer.cpp + DifferentiabilityWitnessDevirtualizer.cpp GenericSpecializer.cpp MergeCondFail.cpp Outliner.cpp diff --git a/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp b/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp new file mode 100644 index 0000000000000..5d7041e302ac7 --- /dev/null +++ b/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp @@ -0,0 +1,72 @@ +//===--- DifferentiabilityWitnessDevirtualizer.cpp ------------------------===// +// +// This source file is part of the Swift.org open source project +// +// Copyright (c) 2020 Apple Inc. and the Swift project authors +// Licensed under Apache License v2.0 with Runtime Library Exception +// +// See https://swift.org/LICENSE.txt for license information +// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors +// +//===----------------------------------------------------------------------===// +// +// Devirtualized differentiability witnesses whose bodies are availabe, by +// turning "differentiability_witness_function" instructions into "function_ref" +// instructions referencing the appropriate function. +// +//===----------------------------------------------------------------------===// + +#include "swift/SIL/SILBuilder.h" +#include "swift/SIL/SILFunction.h" +#include "swift/SIL/SILInstruction.h" +#include "swift/SILOptimizer/PassManager/Transforms.h" + +using namespace swift; + +namespace { +class DifferentiabilityWitnessDevirtualizer : public SILFunctionTransform { + + /// Returns true if and changes were made. + bool devirtualizeDifferentiabilityWitnessesInFunction(SILFunction &f); + + /// The entry point to the transformation. + void run() override { + if (devirtualizeDifferentiabilityWitnessesInFunction(*getFunction())) + invalidateAnalysis(SILAnalysis::InvalidationKind::CallsAndInstructions); + } +}; +} // end anonymous namespace + +bool DifferentiabilityWitnessDevirtualizer:: + devirtualizeDifferentiabilityWitnessesInFunction(SILFunction &f) { + bool changed = false; + llvm::SmallVector insts; + for (auto &bb : f) { + for (auto &inst : bb) { + auto *dfwi = dyn_cast(&inst); + if (!dfwi) + continue; + insts.push_back(dfwi); + } + } + for (auto *inst : insts) { + auto *wit = inst->getWitness(); + if (wit->isDeclaration()) + f.getModule().loadDifferentiabilityWitness(wit); + if (wit->isDeclaration()) + continue; + changed = true; + SILBuilderWithScope builder(inst); + auto kind = inst->getWitnessKind().getAsDerivativeFunctionKind(); + assert(kind.hasValue()); + auto *newInst = builder.createFunctionRefFor(inst->getLoc(), + wit->getDerivative(*kind)); + inst->replaceAllUsesWith(newInst); + inst->getParent()->erase(inst); + } + return changed; +} + +SILTransform *swift::createDifferentiabilityWitnessDevirtualizer() { + return new DifferentiabilityWitnessDevirtualizer(); +} diff --git a/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil b/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil new file mode 100644 index 0000000000000..b13a2b4a4eaa4 --- /dev/null +++ b/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil @@ -0,0 +1,42 @@ +// RUN: %target-sil-opt -differentiability-witness-devirtualizer %s -enable-sil-verify-all | %FileCheck %s + +sil_stage raw + +import Swift +import Builtin + +sil_differentiability_witness [parameters 0] [results 0] @witness_defined_in_module : $@convention(thin) (Float) -> Float { + jvp: @witness_defined_in_module_jvp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) + vjp: @witness_defined_in_module_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) +} + +sil_differentiability_witness [parameters 0] [results 0] @witness_definition_not_available : $@convention(thin) (Float) -> Float + +// This is an example of a witness that is available (via deserialization) +// even though it is not defined in the current module. +// witness for static Swift.Float.+ infix(Swift.Float, Swift.Float) -> Swift.Float +sil_differentiability_witness [parameters 0 1] [results 0] @$sSf1poiyS2f_SftFZ : $@convention(method) (Float, Float, @thin Float.Type) -> Float + +sil @witness_defined_in_module : $@convention(thin) (Float) -> Float + +sil @witness_defined_in_module_jvp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) + +sil @witness_defined_in_module_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) + +sil @witness_definition_not_available : $@convention(thin) (Float) -> Float + +sil public_external [transparent] [serialized] @$sSf1poiyS2f_SftFZ : $@convention(method) (Float, Float, @thin Float.Type) -> Float + +sil @test : $@convention(thin) (Float) -> () { +bb0(%0 : $Float): + %1 = differentiability_witness_function [vjp] [parameters 0] [results 0] @witness_defined_in_module : $@convention(thin) (Float) -> Float + // CHECK: %1 = function_ref @witness_defined_in_module_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float) + + %2 = differentiability_witness_function [vjp] [parameters 0] [results 0] @witness_definition_not_available : $@convention(thin) (Float) -> Float + // CHECK: %2 = differentiability_witness_function [vjp] [parameters 0] [results 0] @witness_definition_not_available : $@convention(thin) (Float) -> Float + + %3 = differentiability_witness_function [vjp] [parameters 0 1] [results 0] @$sSf1poiyS2f_SftFZ : $@convention(method) (Float, Float, @thin Float.Type) -> Float + // CHECK: %3 = function_ref @AD__$sSf1poiyS2f_SftFZ__vjp_src_0_wrt_0_1 : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float)) + + return undef : $() +} From fac70ef460998410e391f4c4863e4f022acd829a Mon Sep 17 00:00:00 2001 From: Ewa Matejska <15254638+ematejska@users.noreply.github.com> Date: Sat, 11 Apr 2020 21:03:06 -0700 Subject: [PATCH 2/3] Fixing the differentiability_witness_inlining.sil test --- .../AutoDiff/SILOptimizer/differentiability_witness_inlining.sil | 1 + 1 file changed, 1 insertion(+) diff --git a/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil b/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil index b13a2b4a4eaa4..3637309bf21c0 100644 --- a/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil +++ b/test/AutoDiff/SILOptimizer/differentiability_witness_inlining.sil @@ -2,6 +2,7 @@ sil_stage raw +import _Differentiation import Swift import Builtin From c641df2a7d150c9a5d8275547650333c69e3da22 Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Sat, 11 Apr 2020 21:47:39 -0700 Subject: [PATCH 3/3] [AutoDiff] NFC: gardening. Fix file header comments. --- .../DifferentiabilityWitnessDevirtualizer.cpp | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp b/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp index 5d7041e302ac7..0ffdc75826a01 100644 --- a/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp +++ b/lib/SILOptimizer/Transforms/DifferentiabilityWitnessDevirtualizer.cpp @@ -10,9 +10,8 @@ // //===----------------------------------------------------------------------===// // -// Devirtualized differentiability witnesses whose bodies are availabe, by -// turning "differentiability_witness_function" instructions into "function_ref" -// instructions referencing the appropriate function. +// Devirtualizes `differentiability_witness_function` instructions into +// `function_ref` instructions for differentiability witness definitions. // //===----------------------------------------------------------------------===// @@ -26,7 +25,7 @@ using namespace swift; namespace { class DifferentiabilityWitnessDevirtualizer : public SILFunctionTransform { - /// Returns true if and changes were made. + /// Returns true if any changes were made. bool devirtualizeDifferentiabilityWitnessesInFunction(SILFunction &f); /// The entry point to the transformation. @@ -50,17 +49,17 @@ bool DifferentiabilityWitnessDevirtualizer:: } } for (auto *inst : insts) { - auto *wit = inst->getWitness(); - if (wit->isDeclaration()) - f.getModule().loadDifferentiabilityWitness(wit); - if (wit->isDeclaration()) + auto *witness = inst->getWitness(); + if (witness->isDeclaration()) + f.getModule().loadDifferentiabilityWitness(witness); + if (witness->isDeclaration()) continue; changed = true; SILBuilderWithScope builder(inst); auto kind = inst->getWitnessKind().getAsDerivativeFunctionKind(); assert(kind.hasValue()); auto *newInst = builder.createFunctionRefFor(inst->getLoc(), - wit->getDerivative(*kind)); + witness->getDerivative(*kind)); inst->replaceAllUsesWith(newInst); inst->getParent()->erase(inst); }