Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SR-9728] [AD] IRGen crasher on a witness thunk when the conforming struct has an @autodiff function stored property #52162

Closed
rxwei opened this issue Jan 23, 2019 · 5 comments
Assignees
Labels
bug A deviation from expected or documented behavior. Also: expected but undesirable behavior. swift for tensorflow

Comments

@rxwei
Copy link
Member

rxwei commented Jan 23, 2019

Previous ID SR-9728
Radar None
Original Reporter @rxwei
Type Bug
Status Resolved
Resolution Done
Additional Detail from JIRA
Votes 0
Component/s Swift for TensorFlow
Labels Bug
Assignee @marcrasi
Priority Medium

md5: 8c3cace3f9ee6e02c7a3ed32a3ce65b2

relates to:

  • SR-9717 [AD] Unhandled use of FunctionRefInst

Issue Description:

import TensorFlow
public protocol Layer: Differentiable & KeyPathIterable
    where AllDifferentiableVariables: KeyPathIterable {
    /// The input type of the layer.
    associatedtype Input: Differentiable
    /// The output type of the layer.
    associatedtype Output: Differentiable

    /// Returns the output obtained from applying to an input.
    @differentiable(wrt: (self, input))
    func applied(to input: Input) -> Output
}

public struct Dense<Scalar>: Layer
    where Scalar : FloatingPoint & Differentiable & TensorFlowScalar {
    @differentiable(vjp: vjpWeight)
    public var weight: Tensor<Scalar>
    @differentiable(vjp: vjpBias)
    public var bias: Tensor<Scalar>
    public typealias Activation = @autodiff (Tensor<Scalar>) -> Tensor<Scalar>
    @noDerivative public let meow: Bool

    // FIXME: Removing this makes compilation succeed.
    @noDerivative public let activation: Activation

    @usableFromInline func vjpWeight() -> (Tensor<Scalar>, (Tensor<Scalar>) -> CotangentVector) {
        return (weight, { v in CotangentVector(weight: v, bias: .zero) })
    }
    @usableFromInline func vjpBias() -> (Tensor<Scalar>, (Tensor<Scalar>) -> CotangentVector) {
        return (bias, { v in CotangentVector(weight: .zero, bias: v) })
    }

    // FIXME
    public var allKeyPaths: [PartialKeyPath<Dense>] {
        return [\Dense.weight, \Dense.bias]
    }

    @differentiable(wrt: (self, input), vjp: _vjpApplied(to:))
    public func applied(to input: Tensor<Scalar>) -> Tensor<Scalar> {
        return matmul(input, weight) + bias
    }

    @usableFromInline
    func _vjpApplied(to input: Tensor<Scalar>)
        -> (Tensor<Scalar>, (Tensor<Scalar>) -> (Dense.CotangentVector, Tensor<Scalar>)) {
      let r0 = matmul(input, weight)
      let r1 = r0 + bias
      let biasShape = bias.shapeTensor
      func pullback(_ v: Tensor<Scalar>) -> (Dense.CotangentVector, Tensor<Scalar>) {
          return (Dense.CotangentVector(
                      weight: matmul(input.transposed(), v),
                      bias: v.unbroadcast(toShape: biasShape)),
                  matmul(v, weight.transposed()))
      }
      return (r1, pullback)
    }
}
Assertion failed: (to <= Values.size()), function getRange, file /Users/rxwei/Development/Swift/swift-tf-source/swift/lib/IRGen/Explosion.h, line 104.
Stack dump:
0.  Program arguments: /Users/rxwei/Development/Swift/swift-tf-source/build/Ninja-ReleaseAssert+stdlib-Release/swift-macosx-x86_64/bin/swift -frontend -c -primary-file test.swift -target x86_64-apple-darwin18.2.0 -enable-objc-interop -sdk /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.14.sdk -color-diagnostics -module-name test -o /var/folders/xz/fy3vsyyn6vs9cqbh32s6z1f400h255/T/test-4fb6d0.o
1.  While emitting IR SIL function "@AD__$s4test5DenseVyxGAA5LayerA2aEP7applied2to6OutputQz5InputQz_tFTW_vjp_SS".
 for 'applied(to:)' at test.swift:49:12
0  swift                    0x0000000105b6d028 llvm::sys::PrintStackTrace(llvm::raw_ostream&) + 40
1  swift                    0x0000000105b6bff8 llvm::sys::RunSignalHandlers() + 248
2  swift                    0x0000000105b6d642 SignalHandler(int) + 258
3  libsystem_platform.dylib 0x00007fff77342b3d _sigtramp + 29
4  libsystem_platform.dylib 0x0000000000003bc8 _sigtramp + 2295074984
5  libsystem_c.dylib        0x00007fff772001c9 abort + 127
6  libsystem_c.dylib        0x00007fff771c8868 basename_r + 0
7  swift                    0x00000001023d5a33 swift::SILInstructionVisitor<(anonymous namespace)::IRGenSILFunction, void>::visit(swift::SILInstruction*) + 60739
8  swift                    0x00000001023c3322 swift::irgen::IRGenModule::emitSILFunction(swift::SILFunction*) + 7938
9  swift                    0x00000001022cbacb swift::irgen::IRGenerator::emitLazyDefinitions() + 715
10 swift                    0x000000010239b176 performIRGeneration(swift::IRGenOptions&, swift::ModuleDecl*, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::StringRef, swift::PrimarySpecificPaths const&, llvm::LLVMContext&, swift::SourceFile*, llvm::GlobalVariable**) + 1254
11 swift                    0x000000010239b5f6 swift::performIRGeneration(swift::IRGenOptions&, swift::SourceFile&, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::StringRef, swift::PrimarySpecificPaths const&, llvm::LLVMContext&, llvm::GlobalVariable**) + 86
12 swift                    0x00000001022228cb performCompile(swift::CompilerInstance&, swift::CompilerInvocation&, llvm::ArrayRef<char const*>, int&, swift::FrontendObserver*, swift::UnifiedStatsReporter*) + 13739
13 swift                    0x000000010221e2b0 swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 2848
14 swift                    0x00000001021d008e main + 686
15 libdyld.dylib            0x00007fff77157ed9 start + 1
16 libdyld.dylib            0x000000000000000f start + 2297069879
<unknown>:0: error: unable to execute command: Abort trap: 6
@rxwei
Copy link
Member Author

rxwei commented Jan 23, 2019

The code currently triggers https://bugs.swift.org/browse/SR-9717, which was fixed by #22054 When that patch gets merged, this will be reproducible from HEAD.

@marcrasi
Copy link
Mannequin

marcrasi mannequin commented Jan 23, 2019

I made a bit of progress on this.

Firstly, here's a reduced reproducer:

import TensorFlow

public protocol Layer : Differentiable {
  associatedtype T: FloatingPoint & Differentiable & TensorFlowScalar
  @differentiable(wrt: (self, input))
  func applied(to input: Tensor<T>) -> Tensor<T>
}

public struct Dense<T : FloatingPoint & Differentiable & TensorFlowScalar> : Layer {
  public var weight: Tensor<T>

  // Compiles successfully if you comment this out.
  @noDerivative public let activation: @autodiff (Tensor<T>) -> Tensor<T>

  @differentiable(wrt: (self, input), vjp: _vjpApplied)
  public func applied(to input: Tensor<T>) -> Tensor<T> {
    return weight * input
  }

  public func _vjpApplied(to input: Tensor<T>) -> (Tensor<T>, (Tensor<T>) -> (Dense.CotangentVector, Tensor<T>)) {
    fatalError("not implemented")
  }
}

Secondly, the problem seems to be that the Differentiation pass is not adding associated functions to the `autodiff_function` instructions inside the protocol witness for `applied(to🙂`. It's very weird that adding `@noDerivative public let activation: @autodiff (Tensor<T>) -> Tensor<T>` causes this to happen.

I'm now going to look at why the Differentitation pass isn't adding the associated functions.

@marcrasi
Copy link
Mannequin

marcrasi mannequin commented Jan 23, 2019

Actually, the Differentiation pass is adding the associated functions to the autodiff_function instructions. But the associated functions are gone when it gets to IRGen. Maybe there is some pass that accidentally removes the associated functions from the autodiff_function instructions? Looking for that now.

@marcrasi
Copy link
Mannequin

marcrasi mannequin commented Jan 24, 2019

Here is a fix: marcrasi@9117ced

It goes on top of #22054 .

Could you just patch that fix onto your PR?

@rxwei
Copy link
Member Author

rxwei commented Jan 24, 2019

Done patching. My bad omitting associated functions! Thanks for the fix.

@swift-ci swift-ci transferred this issue from apple/swift-issues Apr 25, 2022
This issue was closed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug A deviation from expected or documented behavior. Also: expected but undesirable behavior. swift for tensorflow
Projects
None yet
Development

No branches or pull requests

1 participant