From 2ee5bdf741ae14d6f3013d052f5c527aa5a147dc Mon Sep 17 00:00:00 2001 From: Dan Zheng Date: Wed, 12 Feb 2020 15:17:09 -0800 Subject: [PATCH] [AutoDiff] Fix `differentiable_function`-related specialization crashes. Fix crashes related to generic specialization of `partial_apply` operands to `differentiable_function` instructions. `differentiable_function` requires derivative function operand types to match expected derivative function types computed from the original function operand's type, so operands cannot be specialized individually without specializing the others. Resolves TF-891 and TF-1126. --- lib/SILOptimizer/IPO/CapturePropagation.cpp | 14 ++++- lib/SILOptimizer/Utils/Generics.cpp | 6 +- ...26-derivative-generic-specialization.swift | 56 +++++++++++++++++++ ...891-protocol-req-capture-propagation.swift | 6 +- 4 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 test/AutoDiff/downstream/compiler_crashers_fixed/tf1126-derivative-generic-specialization.swift rename test/AutoDiff/downstream/{compiler_crashers => compiler_crashers_fixed}/tf891-protocol-req-capture-propagation.swift (97%) diff --git a/lib/SILOptimizer/IPO/CapturePropagation.cpp b/lib/SILOptimizer/IPO/CapturePropagation.cpp index 4294f5e3f8d72..3e0c6a273a195 100644 --- a/lib/SILOptimizer/IPO/CapturePropagation.cpp +++ b/lib/SILOptimizer/IPO/CapturePropagation.cpp @@ -408,6 +408,17 @@ static SILFunction *getSpecializedWithDeadParams( return nullptr; } + // SWIFT_ENABLE_TENSORFLOW + // Disable specialization for instructions that are operands of + // `differentiable_function` instructions. `differentiable_function` + // requires derivative function operand types to match expected derivative + // function types computed from the original function operand's type, so + // operands cannot be specialized individually without specializing the + // others. + if (!PAI->getUsersOfType().empty()) + return nullptr; + // SWIFT_ENABLE_TENSORFLOW END + auto Rep = Specialized->getLoweredFunctionType()->getRepresentation(); if (getSILFunctionLanguage(Rep) != SILFunctionLanguage::Swift) return nullptr; @@ -463,7 +474,7 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) { if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder, PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) { // SWIFT_ENABLE_TENSORFLOW - // Add a previously unexercised check to prevent AD crash. Rewrite + // Add a previously unexercised check to prevent AutoDiff crash. Rewrite // `partial_apply` only if the specialized function is `@convention(thin)`. // Revert check when `VJPEmitter::visitApplyInst` no longer produces // argumentless `partial_apply` instructions. @@ -476,6 +487,7 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) { } return true; } + // SWIFT_ENABLE_TENSORFLOW END } // Second possibility: Are all partially applied arguments constant? diff --git a/lib/SILOptimizer/Utils/Generics.cpp b/lib/SILOptimizer/Utils/Generics.cpp index 91152e4274242..f2964e0bb88a4 100644 --- a/lib/SILOptimizer/Utils/Generics.cpp +++ b/lib/SILOptimizer/Utils/Generics.cpp @@ -531,7 +531,11 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee, // SWIFT_ENABLE_TENSORFLOW // Disable specialization for instructions that are operands of - // `differentiable_function` instructions. + // `differentiable_function` instructions. `differentiable_function` + // requires derivative function operand types to match expected derivative + // function types computed from the original function operand's type, so + // operands cannot be specialized individually without specializing the + // others. if (Apply.getInstruction()) for (auto result : Apply.getInstruction()->getResults()) for (auto use : result->getUses()) diff --git a/test/AutoDiff/downstream/compiler_crashers_fixed/tf1126-derivative-generic-specialization.swift b/test/AutoDiff/downstream/compiler_crashers_fixed/tf1126-derivative-generic-specialization.swift new file mode 100644 index 0000000000000..ad6d462c7e4ce --- /dev/null +++ b/test/AutoDiff/downstream/compiler_crashers_fixed/tf1126-derivative-generic-specialization.swift @@ -0,0 +1,56 @@ +// RUN: %target-swift-frontend -O -emit-sil %s -verify +// REQUIRES: asserts + +// TF-1126: Generic specialization crash during capture propagation. +// Related to `@differentiable` function with `partial_apply` operands, +// to be specialized. Occurs only with `-O`. + +struct A: Differentiable{ + var b: SIMD8 +} + +@differentiable +func function(a: A) -> A { + var a = a + a.b = a.b - SIMD8(repeating: 1.0) + return a +} + +let masks: [SIMD8] = [[1,0,0,0,0,0,0,0], + [0,1,0,0,0,0,0,0], + [0,0,1,0,0,0,0,0], + [0,0,0,1,0,0,0,0], + [0,0,0,0,1,0,0,0], + [0,0,0,0,0,1,0,0], + [0,0,0,0,0,0,1,0], + [0,0,0,0,0,0,0,1]] + +extension SIMD8 where Scalar == Float{ + @differentiable(where Scalar: Differentiable) + func updated(at index: Int, with newValue: Scalar) -> Self { + let mask = masks[index] + let result = self - (self * mask) + (newValue * mask) + return result + } +} + +// Looking for a function: $ss4SIMDPss14DifferentiableRzSB6Scalars11SIMDStoragePRpzsAA13TangentVectorsACPRpzSBAhI_AdFRPzrlE12_vjpSubtract3lhs3rhsx5value_AJ_AJtAJc8pullbacktx_xtFZs5SIMD8VySfG_Tg5 +// Expected type: @convention(method) (@in_guaranteed SIMD8, @in_guaranteed SIMD8, @thick SIMD8.Type) -> (@out SIMD8, @owned @callee_guaranteed (@in_guaranteed SIMD8) -> (@out SIMD8, @out SIMD8)) +// Found type: @convention(method) (SIMD8, SIMD8, @thick SIMD8.Type) -> (@out SIMD8, @owned @callee_guaranteed (@in_guaranteed SIMD8) -> (@out SIMD8, @out SIMD8)) +// Assertion failed: (ReInfo.getSpecializedType() == SpecializedF->getLoweredFunctionType() && "Previously specialized function does not match expected type."), function lookupSpecialization, file /Users/swiftninjas/s4tf/swift/lib/SILOptimizer/Utils/Generics.cpp, line 1833. +// Stack dump: +// ... +// 1. Swift version 5.2-dev (Swift bf631dc2e4) +// 2. While running pass #113021 SILFunctionTransform "CapturePropagation" on SILFunction "@AD__$ss5SIMD8V6deleteSfRszrlE7updated2at4withABySfGSi_SftF__vjp_src_0_wrt_1_2". +// for 'updated(at:with:)' (at /Users/porter/Dropbox (PassiveLogic)/Team/Team Members Scratch Space/Porter/Experiments/Playgrounds/delete/delete/main.swift:75:5) +// llvm::sys::PrintStackTrace(llvm::raw_ostream&) + 37 +// llvm::sys::RunSignalHandlers() + 85 +// SignalHandler(int) + 278 +// _sigtramp + 29 +// _sigtramp + 2821162056 +// abort + 127 +// basename_r + 0 +// swift::GenericFuncSpecializer::lookupSpecialization() (.cold.1) + 35 +// swift::GenericFuncSpecializer::lookupSpecialization() + 2109 +// (anonymous namespace)::CapturePropagation::optimizePartialApply(swift::PartialApplyInst*) + 1301 +// (anonymous namespace)::CapturePropagation::run() + 265 diff --git a/test/AutoDiff/downstream/compiler_crashers/tf891-protocol-req-capture-propagation.swift b/test/AutoDiff/downstream/compiler_crashers_fixed/tf891-protocol-req-capture-propagation.swift similarity index 97% rename from test/AutoDiff/downstream/compiler_crashers/tf891-protocol-req-capture-propagation.swift rename to test/AutoDiff/downstream/compiler_crashers_fixed/tf891-protocol-req-capture-propagation.swift index f146c02c5f7e9..37c8dd769d2ea 100644 --- a/test/AutoDiff/downstream/compiler_crashers/tf891-protocol-req-capture-propagation.swift +++ b/test/AutoDiff/downstream/compiler_crashers_fixed/tf891-protocol-req-capture-propagation.swift @@ -1,6 +1,10 @@ -// RUN: not --crash %target-swift-frontend -O -emit-ir %s +// RUN: %target-swift-frontend -O -emit-ir %s // REQUIRES: asserts +// TF-891: Generic specialization crash during capture propagation. +// Related to `@differentiable` function with `partial_apply` operands, +// to be specialized. Occurs only with `-O`. + public protocol Protocol: Differentiable { @differentiable func requirement1(_ arg: T) -> Float