[SR-15808] [AutoDiff] Non-differentiable closure parameters not parsed correctly

[philipturner]

Previous ID	SR-15808
Radar	None
Original Reporter	@philipturner
Type	Bug

Attachment: Download

Additional Detail from JIRA

Votes	0
Component/s
Labels	Bug
Assignee	None
Priority	Medium

md5: a3df518fba66afe2f96d3bbdb0bf649a

Issue Description:

The following code produces a compiler crash during SILGen because it skipped classifying the closure as <<error type>> during AST creation:

import _Differentiation

protocol DifferentiableCollection: Sequence {}

extension DifferentiableCollection {
  func differentiableMap(_ body: @differentiable (Void) -> Void) {
    fatalError()
  }
}

While this code, which correctly identifies the closure as <<error type>> during AST creation, does not crash the compiler:

import _Differentiation

protocol DifferentiableCollection: Sequence {}

extension DifferentiableCollection {
  func differentiableMap(_ body: @differentiable (Int) -> Void) {
    fatalError()
  }
}

The crash log is attached.

[AnthonyLatsis]

Have you already checked whether `TypeChecker::isDifferentiable` returns true when passed Void (which it shouldn't IIUC)?

[philipturner]

I have narrowed down the bug's source to line 594 on TypeChecker.cpp. isDifferentiable should (theoretically) return false for Void. It does return false for Int. But, it actually returns true for Void.

I need to ensure that we can reproduce this behavior for any arbitrary Void, not just one found in this bug's peculiar context. I hypothesize that the bug is called by the isDifferentiable function being messed up, always classifying Void as differentiable.

[philipturner]

@AnthonyLatsis I don't have the experience necessary to test the hypothesis above efficiently. Do you have the resources to do that at the moment?

[AnthonyLatsis]

It does indeed return true for Void. But it doesn't feel like an oversight. Are you sure that Void is not supposed to be differentiable? It looks like tuples are (if their elements are), and Void is just a special case – the empty tuple.

[philipturner]

But the empty tuple doesn't have any differentiable elements! They might have set it to be differentiable because it's the return type of pullbacks in mutating functions. Could you find a way to conditionally set empty tuples to not-differentiable, and see what happens?

[philipturner]

Or we could just make an exception at the problematic call site, where we bypass this if it's an empty tuple. If you change it to register an empty tuple as non-differentiable, I might ping some AutoDiff experts for their insight.

[AnthonyLatsis]

I can, but I'd be taking shots in the dark. If Void is supposed to be differentiable only in certain structural positions, we need to be certain about the criteria. But if the code happens to be deliberate, the bug lies somewhere else.

[philipturner]

I don't think this behavior is intentional. Void should not be differentiable, and in some examples below, the system gives diagnostics about that. When you pass a function in as a closure parameter, then it starts behaving differently than normal.

The following crash:

func helloWorld(_ x: @differentiable (Void) -> Void) {}
func helloWorld(_ x: @differentiable (Void) -> Float) {}
func helloWorld(_ x: @differentiable (Float) -> Void) {}
func helloWorld(_ x: @differentiable (@noDerivative Float, Void) -> Float) {}

The following produce an error diagnostic:

// These functions are inconsistent with their counterparts above
@differentiable func helloWorld(_ x: Void) -> Void {}
@differentiable func helloWorld(_ x: Void) -> Float {}
@differentiable func helloWorld(_ x: Float) -> Void {}
func helloWorld(_ x: @differentiable (@noDerivative Float) -> Float) {}

func helloWorld(_ x: @differentiable (@noDerivative Void) -> Void) {}
func helloWorld(_ x: @differentiable (@noDerivative Void) -> Float) {}

@differentiable func helloWorld(_ x: inout Float, y: Float) -> Int {}
@differentiable func helloWorld(_ x: inout Float, y: Float) -> (Void, Int) {}
@differentiable func helloWorld(_ x: inout Float, y: Float) -> (Void, Void) {}

func helloWorld(_ x: @differentiable (Float) -> Int) {}
func helloWorld(_ x: @differentiable (inout Float, Float) -> Int) {}
// Counterpart cited below
func helloWorld(_ x: @differentiable (inout Float, Float, Int) -> Void) {}

The following do not produce an error diagnostic or crash:

@differentiable func helloWorld(_ x: inout Float, y: Float) -> Void {}
@differentiable func helloWorld(_ x: inout Float, y: Float) -> (Void) {}
@differentiable func helloWorld(_ x: inout Float, y: Float, z: Void) -> (Void) {}
// This function is inconsistent with its counterpart above
@differentiable func helloWorld(_ x: inout Float, y: Float, z: Int) -> (Void) {}

func helloWorld(_ x: @differentiable (inout Float, Float) -> Void) {}
func helloWorld(_ x: @differentiable (inout Float, Float) -> (Void)) {}
func helloWorld(_ x: @differentiable (inout Float, Float) -> ((Void))) {}
func helloWorld(_ x: @differentiable (inout Float, Float) -> (((Void)))) {}

This means Swift has special treatment for when a function's return is either Void or recursive tuple of Void. It should never classify Void as differentiable in the first place; rather, it should do special-case checking at the return site of a function. @AnthonyLatsis how did you prove that Swift always says Void is differentiable? The diagnostics above suggest otherwise.

Some interesting combinations from above:

// This compiles successfully
@differentiable func helloWorld(_ x: inout Float, _ y: Float, _ z: Int) -> Void {}

// This produces an error
func helloWorld(_ x: @differentiable (inout Float, Float, Int) -> Void) {}

// These both compile successfully
@differentiable func helloWorld(_ x: inout Float, _ y: Float, _ z: Void) -> (Void) {}
func helloWorld(_ x: @differentiable (inout Float, Float, Void) -> Void) {}

// These both compile successfully
@differentiable func helloWorld(_ x: inout Float, _ y: Float, _ z: @noDerivative Int) -> (Void) {}
func helloWorld(_ x: @differentiable (inout Float, Float, @noDerivative Int) -> Void) {}

// These both compile successfully
@differentiable func helloWorld(_ x: inout Float, _ y: Float, _ z: @noDerivative Void) -> (Void) {}
func helloWorld(_ x: @differentiable (inout Float, Float, @noDerivative Void) -> Void) {}

[philipturner]

I think the compiler is supposed to insert an implicit `@noDerivative` at some point in AST. If so, and it was just skipping that step in some edge cases, that would explain everything.

Except this crasher:

func helloWorld(_ x: @differentiable (Float) -> (Void)) {}

Its assertion failure crash says "result indices must not be empty". This parallels the original crasher that started this thread, where instead parameter indices were empty.

[philipturner]

Fix submitted as #41174