Store `@differentiable` attribute requirements as `GenericSignature` instead of `ArrayRef<Requirement>`s?

Development Compiler
1

Hi compiler experts,

I’m working on the differentiable programming project and have a question about the best representation for where clause requirements.

The @differentiable attribute marks function-like declarations as being differentiable. On generic declarations, the @differentiable attribute may have an associated where clause with requirements for conditional differentiability. Examples:

@differentiable(where T: Differentiable)
func identity<T>(_ x: T) -> T { x }

struct Tensor<Scalar> { ... }
extension Tensor where Scalar: Numeric {
  // Numeric operations are differentiable only if the scalar type is
  // floating-point (continuous) and differentiable.
  @differentiable(where Scalar: FloatingPoint & Differentiable)
  static func +(lhs: Self, rhs: Self) -> Self { ... }
}

So far, we’ve represented these generic requirements as ArrayRef<Requirement>, stored and serialized in the AST DifferentiableAttr and the SILDifferentiableAttr data structures. However, a common operation is to compute a “derivative generic signature” given the original declaration’s generic signature and the attribute requirements: see getAutoDiffAssociatedFunctionGenericSignature. This code is duplicated in 2-3 places.

I wonder if it makes more sense to store requirements directly as a GenericSignature or GenericEnvironment in DifferentiableAttr? This should eliminate redundant code for computing derivative generic signatures from ArrayRef<Requirement> and the original declaration GenericSignature. There is no need to distinguish original requirements vs attribute requirements.

Currently, the @differentiable type-checking code actually already forms a derivative GenericSignature and GenericEnvironment when resolving RequirementReprs during type-checking, so we can just store one of those instead of the resolved Requirements.

I also wonder whether it's appropriate to store a GenericSignature or GenericEnvironment? The primary use for the generic signature/environment is to remap types, which seems easier with GenericEnvironment but maybe possible with both:

  • GenericEnvironment: GenericEnvironment::mapTypeIntoContext. Seems robust.
  • GenericSignature: no dedicated method for remapping types. Possible to use Type::subst(SubstitutionMap) with GenericSignature::getIdentitySubstitutionMap but seems less robust.

cc @Slava_Pestov

1 Like
2

Yeah, it makes sense to at least extract this into a utility function, or to store the resulting generic signature somewhere.

It would help if you clarified what you're actually doing by remapping types. The two functions you mention are different -- mapTypeIntoContext() replaces generic parameters with the archetypes from the generic environment. If you're emitting a SIL function whose generic signature is the derived generic signature with the additonal requirements, you probably want to set the SILFunction's generic environment to be the derived generic environment, and use mapTypeIntoContext() to create types for use in SIL instructions in the body.

On the other hand, calling subst() with getIdentitySubstitutionMap() should be a no-op if the original type used type parameter types already. If the original type used archetypes from a generic environment built from the signature, then substituting with the identity map will produce a type written with type parameters again, but Type::mapTypeOutOfContext() is a more idiomatic way of achieving that.

A generic signature itself does not have archetypes associated with it -- that's why you can create zero or more generic environments from a generic signature. We plan on simplifying this at some point so that each signature has exactly one generic environment which is lazily built as needed. Once this is done you can think of the two concepts are mostly interchangeable.

1 Like
3

Yes, our use case is exactly "creating a SIL function with the derivative generic environment with additional requirements, and using the environment to remap types via TypeSubstCloner machinery".

Thanks for this clarification, I’ve been wondering about this distinction. Generic environments have archetypes while generic signatures do not (only generic type parameter types and requirements).

Tentatively, we’ll pursue storing a derivative GenericSignature in AST DifferentiableAttr and SILDifferentiableAttr then. Directly storing GenericSignature will reduce code duplication and fix issues where derivative generic signature is currently computed incorrectly. Derivative SIL functions will have an (optional) generic environment created from this GenericSignature, same as today.

Eventually, differentiability witness tables will replace SILDifferentiableAttr and store an optional derivative GenericSignature per entry.

Thanks again!

1 Like
4

Bump: (SIL)SpecializeAttr have been changed to store GenericSignature instead of ArrayRef<Requirement> in this commit.

This validates the approach for (SIL)DifferentiableAttr. @marcrasi found that storing rather than recomputing GenericSignature should significantly improve compiler performance (during type-checking, SILGen, differentiation transform).

1 Like