Overhaul user-defined override system for higher-order operators (e.g. `fwd_diff`)

[saipraveenb25]

Our higher-order operators such as fwd_diff and bwd_diff allow user-defined overrides on a per-function basis.

E.g. The [BackwardDerivative(deriv_fn)] decorator on fn says that bwd_diff(fn) should simply resolve to deriv_fn

For generically defined functions, we have a concept of "sub-type functions".
E.g.: A definition such as T f<T>(T x) { } works with any T, while T f<T : IFloat>(T x) works on a strict subset of possible specializations of the original 'f'.
We would like to be able to provide derivative overrides that apply to a sub-set of possible specializations of a function f

This needs some re-factoring of the current system where we associate each declaration with its derivative declaration.

We will instead need to construct this mapping at the relevant use-site of the declaration, and use the specialization arguments to determine which override applies.

[saipraveenb25]

OUTDATED PROPOSAL (See shader-slang/spec#11 for latest)

Summary of Planned Changes:

Parser

Function Operator Declarations

__function_operator operator_name : operator_class; 

where operator_class is a subtype of FunctionOperatorDecl.

Should be parsed as:

    operator_class (name=operator_name, syntaxExpr=expr)

Parsing higher order operator invocations

Allowable syntax:

operator_name(expr)(args...)

Should be parsed as:

    InvokeExpr(HigherOrderOperatorExpr(operatorDecl, expr), args)

Operator overloads for (func -> func) operators (e.g. ForwardDerivative)

Allow for the following syntax for function operator overloads:

    __generic<generic_param_decls>
    ReturnType func_name<specializations(generic_param_decls)>.operator operator_name(param_types(generic_param_decls))
    {  /* ... */  }

Should be parsed as:

GenericDecl
{
    ParamDecl (generic_params)
    FunctionOperatorOverloadDecl (operatorDecl, fnExpr)
    FuncDecl // operator function is nameless
    {
        ParamDecl param_types(generic_params)
    }
}

For the first step, we'll continue to use the override syntax that uses decorators (i.e. [ForwardDerivativeOf(fn)]), though we do need to upgrade the syntax to allow for new operators

Operator overloads for (func -> type) operators (e.g. IntermediateContextType)

Allow for the following syntax for function operator overloads (that return a type):

    __generic<generic_param_decls>
    func_name<specializations(generic_param_decls)>.operator operator_name = TypeExpr(generic_param_decls);

Semantics

From now on, the link between a function and its higher order operator overloads should come from the overload to the function, and not the other way around. Previously, the 'main' way to create a link was from the function to the overload.
This new system makes it much more like operator overloading, and allows overloads to apply to a sub-set of specializations of the target functions (like extensions).

Checking of [...DerivativeOf(fnExpr)] attributes (or FunctionOperatorOverloadDecl(fnExpr))

This checking is a bit more complex that other inference areas, but we can use the same trySolveConstraintSystem() mechanism used elsewhere in the codebase. It's just that the constraints come from a different place: the argument types + any existing substitutions on the function's decl-ref, so we can't directly use inferGenericArguments() to solve for the generic parameters.

• Check fnExpr
• If fnExpr is a DeclRefExpr, then just register the declaration with the appropriate decl under the appropriate operator decl, so that it can be retreived from getOperatorImplementationsForDecl(decl, operatorDecl), then we only have one candidate.
• If fnExpr is an OverloadedExpr, then we will go through each candidate and see which one works.
• For each candidate, we'll apply the following logic:
  • If candidate is already a FunctionDeclBase, then we'll first calculate the type of this from the function's context, and use trySolveConstraintSystem.
    to construct a reference to overloadThisType from unifying operatorDecl.translateThisType(funcThisType) with overloadThisType (using the funcDeclRef's thisType as the base). If the unification fails, skip this candidate.
  • Then, we'll apply trySolveConstraintSystem to construct a DeclRef(overloadDecl) from the candidate funcDeclRef's context, by unifying the funcDeclRef's parameter types with the overloadDecl's default parameter types. We will also add the constraints from the funcDeclRef's substitutions. We should end up with a fully specialized overloadDeclRef. If this fails, skip this candidate.
  • Finally, we'll construct a InvokeExpr(overloadDeclRef, translatedArgs), and check this. If this works, then we've found a valid candidate.
  • If the original candidate was a GenericDecl, then we do all the same things, except skip adding the substitutions when inferring generic arguments (since there aren't any).
  • If it's a PartiallyAppliedGenericExpr, then we do all the same things, but add in the knownGenericArgs as constraints, when performing the inference.
  • Anything else is an error.
• If we find no candidate, it's an error (no matching function found).
• If we find multiple candidates, it's an error (ambiguous).
• Otherwise, we'll add the candidate to the appropriate resolved function's decl's operator implementations list. Note, we only store a decl and not a decl-ref, so thing need to get re-resolved at the call-site.

HigherOrderOperatorExpr Base class for higher-order operator expressions.

• Has an expr field which is the operand of the operator.
• Has an opDecl field which is SyntaxClass<FunctionOperatorDecl> to reference the higher-order operator
  declaration.
• Has an isResultFuncType field which is true if the operator returns a function, and false otherwise (returns a type).

FunctionOperatorDecl Base class for higher-order operator declarations.

• Has a translateFuncType(FuncType*) -> FuncType* method that can be used to translate the function type to the resulting type after applying the operator. FuncType here only has paramTypes and resultType. Note that all the types may be ModifiedTypes to indicate modifiers like inout, ref, constref, etc..
  • Unimplemented if the operator returns a type rather than a function.
• Has a translateThisType(Type*) -> Type* method that can be used to translate the this type to the resulting type after applying the operator. This may also be a modified type.
  • Unimplemented if the operator returns a type rather than a function.
• Has a tryGetSynthesisWitness(SemanticsVisitor*, DeclRef*) -> FunctionOperatorWitness* method that can be used to synthesize a witness for the operator.

Requirements and Witness for higher-order operators.

• A function operator can be implicitly treated as a 'requirement' for a function.
  There are 5 such function operators:

• ForwardDerivative

• BackwardDerivative

• BackwardPropagate

• EvalPrimalContext

• IntermediateContextType
  Each of these has a corresponding FunctionOperatorDecl that represents the function operator in the AST.

• A FunctionOperatorWitness is an abstract Witness type that represents how a function satisfies a higher-order operator.

  • It has a opDeclRef field which is a DeclRef to the FunctionOperatorDecl that represents the higher-order operator.
  • It has a funcDeclRef field which is a DeclRef to the base function on which the operator is being applied.
    There are several possible witnesses that represent the different ways a function can satisfy a higher-order operator.

• DeclaredImplementationWitness(satisfyingDeclRef) is created when there is an operator overload definition that applies:

    ReturnType func<spec(generic_params)>.operator<generic_params> operator_name(args(generic_params))
    {  /* ... */  }

  Similarly, for an implementation of an operator that results in a type.

    func<spec(generic_params)>.operator<generic_params> operator_name = TypeExpr;

  DeclaredImplementationWitness can also be used on interface requirements, if the target supports the operator.

    interface IFoo { ResultType bar(Args); }

  In this case, the satisfyingDeclRef

• SynthesizedImplementationWitness(operatorDeclRef, targetFuncDeclRef) is created when operator application can be synthesized from the underlying function through the application of an intrinsic.

  • There are subtypes of such witnesses for each of the function operators (since they all support IR-level synthesis). This witness will lower to a specific IR opcode that invokes the IR translation passes.
  • This witness requires that the function have a body. (i.e. this witness can't be created for intrinsic functions or interface requirements)

• CastToPrimalContextWitness(primalFuncDeclRef) is created when the EvalPrimalContext function can be synthesized from the primal function by wrapping it in a function that accepts an IntermediateContextType parameter.

• CastToBackwardPropagateWitness(backwardDerivativeFuncDeclRef) is created when the BackwardPropagate function can be synthesized from the BackwardDerivative function by wrapping it in a function that accepts an IntermediateContextType parameter.

• ComposeEvalPrimalContextWithPropagateWitness(evalPrimalContextFuncDeclRef, backwardPropagateFuncDeclRef, intermediateContextType) is created when the BackwardDerivative function can be synthesized from the composition of available user-provided EvalPrimalContext and BackwardPropagate functions.

LookupDeclRef-based higher-order operators.

Every higher-order operator has a corresponding Decl node (of parent type FunctionOperatorDecl) that represents the higher-order operator in the AST.

A higher-order applied function can now be represented as LookupDeclRef(operatorDecl, funcDeclRef). A higher-order applied generic can be represented as LookupDeclRef(operatorDecl, genericDeclRef), though this will yield an unusable type until the generic is substituted.

An application of a higher-order operator will first get translated to the appropriate AST node, such as ForwardDerivativeOperator(expr). Upon checking, this node will be processed as follows:

• Check expr.

• If expr is an OverloadedExpr, then we'll simply push the operator into each item, and check each resulting expression. TODO: Can we do any trimming at this point?

• If expr is a DeclRefExpr, we'll try to construct a witness with tryGetFunctionOperatorWitness(declRef, operatorDecl).

  • If this succeeds, we'll return a LookupDeclRef(operatorDecl, declRef, witness).
  • If this fails, then emit a diagnostic at this point "declRef does not support the operator (not eligible for synthesis, and no user-provided definition found)"

• If expr's type is GenericKind, then we'll set the type of ForwardDerivativeOperator to the same GenericKind. This should cover GenericAppExpr and PartiallyAppliedGenericExpr.

• Anything else should produce an error: "unsupported operand type for higher-order operator. expected a callable or a generic function".

• tryResolve() for LookupDeclRef will now also call tryGetFunctionOperatorWitness() each time the witness changes during substitution, and if it is a DeclaredImplementationWitness and refers to a concrete function, we'll resolve it to a DirectDeclRef(declaredDeclRef).

• A LookupDeclRef can only be created with a valid witness that expresses how the given operator is implemented for the given function.

• For cases that are still unresolved, it must continue to be expressed as HigherOrderOperatorExpr(operatorDecl, funcDeclRef) until ResolveInvoke(). If it cannot be resolved to a DeclRef by that point (either directly or through some witness), it is an error.

Overload resolution must be performed with the latter in mind.

Satisfying function operator requirements with tryGetFunctionOperatorWitness(declRef, operatorDecl)

In the new implementation, most of the special case logic is bottlenecked through this function, which is responsible for figuring out if and how a given declRef can be used with operatorDecl.

Logic:

• If getOperatorImplementationsForDecl(declRef.getDecl(), operatorDecl) will return a list of operator overload definitions that reference the given object (this list is propagated across modules similar to how extensions are recorded for declarations).
• For each operatorOverloadDecl, we will use the same logic in applyExtensionToType to check if the given operator overload can be applied to the given declRef. This will need to solve for any generic parameters, and if successful, will return a fully specialized operatorOverloadDeclRef.
• If there's more than one, it's currently an error due to ambiguity.
• If there's exactly one, we'll return a DeclaredImplementationWitness(operatorOverloadDeclRef).
• If there's none, we'll try to synthesize a witness.
  • Since synthesis logic will vary depending on the operator, we'll dispatch on the operator declaration.
    FunctionOperatorDecl::tryGetSynthesisWitness() may provide a witness that satisfies the operator via some intrinsic (that is handled in the backend).
    E.g. ForwardDerivativeOperatorDecl::tryGetSynthesisWitness() will synthesize a SynthesizedImplementationWitness(forwardDerivativeOperatorDecl, declRef), but only if the
    the referenced function has a body (i.e. not an intrinsic or interface requirement), and satisfies other requirements for synthesis (e.g. [Differentiable] attribute).
  • Note that this is where the PrimalSubstitute operator comes into play. FunctionOperatorDecl::tryGetSynthesisWitness() could call tryGetFunctionOperatorWitness(declRef, primalSubstituteOperatorDecl) to locate a substitute declaration if necessary, and then use that in the SynthesizedImplementationWitness to direct the backend to use that function.
  • Similarly, other composite operations like ComposeEvalPrimalContextWithPropagateWitness can be generated by trying to build witnesses for the components.
• This logic requires circularity detection, since one tryGetFunctionOperatorWitness may call another, which ends up calling the first again. If this happens, we should just gracefully return nothing, so that the caller can move on to the next possibility.

AddOverloadCandidate() logic

• We now check if there is a HigherOrderOperatorExpr() expression in the lookup context.

  • If we get this expression, it means we haven't resolved it into a LookupDeclRef yet (in the latter case, it will be resolved as a DeclRefExpr).
  • Usually, this is because we haven't specialized the inner operartor yet! (There may be un-specialized generic parameters in both the operator ad the function).
  • We'll create a tryResolveHigherOrderOperator() function that takes the generic context, and attempts to solve for the generic parameters, by continuously unwrapping the HigherOrderOperatorExpr until we get a generic decl, at each step applying translateFuncType to get the expected type of the arguments in terms of the inner decl's type parameters. Then we inferGenericParameters to get the substituted decl-ref.
  • Then, we re-wrap the expr in a new set of HigherOrderOperatorExprs and then try to resolve the expression.
  • If it's still a HigherOrderOperatorExpr, then we failed. Skip.
  • Otherwise, run it through AddDeclRefCandidate()

• In AddDeclRefCandidate(),

  • It is possible that the result is already a regular function or type decl-ref, if the lookup resolved to a specific function or type. In this case, there's nothing to change.
  • Otherwise, we may have a LookupDeclRef to a FunctionOperatorDecl, which is basically a function that will be synthesized at some point,
    so there's no body for AddFuncOverloadCandidate() to work with, instead it will go through the AddFuncExprCandidate() path, that checks just the type.
  • For now, this is fine, but normally, it's a good idea to make sure the usual AddFuncOverloadCandidate() path works just fine with only the FuncType

Checking of [Differentiable] function bodies.

• We currently override CheckTerm() and register the Type* of a checked term in a dictionary on the attribute, if the type implements IDifferentiable.
• We need to change this logic to also check any DeclRefExpr* in the differentiable function body, then construct higher-order operator expressions for the 5 requirements, check them, and store the results in another dictionary.

IR Lowering

Lowering of [Differentiable] function bodies.

• Register a 'type-hook' that is called whenever a Type* is emitting into the function body. Emit IRDifferentiableTypeAnnotation instructions in the IR for any term that has a Type* registered (in the attribute's dictionary).
• Register a 'decl-ref-hook' that is called whenever a DeclRefExpr* is emitting into the function body. Emit IRAssociatedFuncAnnotation instructions in the IR for any term that has a DeclRefExpr* registered (in the attribute's dictionary).
• IRDifferentiableTypeAnnotation and IRAssociatedFuncAnnotation are both HOISTABLE.
• IRDifferentiableTypeAnnotation and IRAssociatedFuncAnnotation both derive from IRAnnotationBase

Lowering of [Differentiable] function signatures (in interfaces)

• A [Differentiable] function signature in an interface will lower all 5 requirements as additional requirements on the interface.
• This is handled by constructing the DeclaredImplementationWitness for each operator requirement and lowering it to an interface reqiurement entry (with an associated key),
• During lowering of a declRef, in the logic to lower an interface requirement decl, if the target is a FunctionOperatorDecl, we lower it as an extra interface requirement by adding a getFunctionOperatorInterfaceRequirementKey(context, decl, operatorDecl) that returns a key to use for the interface requirement.

Lowering of the new witnesses.

• SynthesizedImplementationWitness is lowered depending on the sub-class.

  • SynthesizedForwardDerivativeWitness is lowered to IRForwardDifferentiate(funcIRValue)
  • SynthesizedBackwardDerivativeWitness is lowered to IRComposeContextAndPropagateFuncs(IRExtractPrimalContextFunc(IRBackwardDifferentiate(funcIRValue)), IRExtractBackwardPropagateFunc(IRBackwardDifferentiate(funcIRValue))).
  • SynthesizedBackwardPropagateWitness is lowered to IRExtractBackwardPropagateFunc(IRBackwardDifferentiate(funcIRValue))
  • SynthesizedEvalPrimalContextWitness is lowered to IRExtractPrimalContextFunc(IRBackwardDifferentiate(funcIRValue))
  • SynthesizedIntermediateContextTypeWitness is lowered to IRExtractIntermediateContextType(IRBackwardDifferentiate(funcIRValue))

• CastToPrimalContextWitness is lowered to IRCastToPrimalContextFunc(funcIRValue)

• CastToBackwardPropagateWitness is lowered to IRCastToBackwardPropagateFunc(funcIRValue)

• ComposeEvalPrimalContextWithPropagateWitness is lowered to IRComposeContextAndPropagateFuncs(evalPrimalContextFuncIRValue, backwardPropagateFuncIRValue)

• If we see a DeclaredImplementationWitness, then the target must be an interface requirement (or something abstract like that). There's nothing to do here. This needs to get handled with performing emitDeclRef(LookupDeclRef(interfaceReqDeclRef, operatorDecl)).

• All translation IR ops are HOISTABLE.

• All translation IR ops derive from IRTranslateBase.

Lowering of FunctionOperatorDecl

The IR form of the operator decl is mostly just to use as a key for various annotations (and in the future, to use as witness table entry keys).
A new IRFunctionOperator global instruction can be used to represent the operator.

• These need to be removed after the translation passes are complete.

Backend Translation Passes

Changes to backend emit logic.

The transcriber passes can now be refactoried into much simpler passes that just service the translation IR ops.

TranslatePass: Main translation pass.

Each IR op-code has a translator class in the backend.
Our main translation pass tries to process each IR translation request in the following manner:

• Gather all IRTranslateBase instructions in the current scope.
• Remove any instruction which does not have a non-annotation use (unwrapping any other IRTranslateBase or IRSpecialize nodes in the process).
• For each remaining instruction, dispatch the appropriate translator to process the translate op.

TranslatorBase: Base class for translators.

• canTranslateInst(IRTranslateBase*) -> bool: Returns true if the op is ready for translation (may depend on other instructions in the same scope).

• processTranslateInst(IRTranslateBase*) -> IRInst*: Accepts the appropriate translate op, and returns a constructed IRInst*.

• E.g.: ForwardDerivativeTranslator : TranslatorBase processes IRForwardDifferentiate and returns a IRFunc (or IRSpecialize of IRFunc).

• E.g.: BackwardDerivativeTranslator : TranslatorBase processes IRBackwardDifferentiate and returns a IRBackwardDerivativeResultTuple(IRFunc, IRFunc, IRType) which contains the primal, adjoint and intermediate context type respectively.

• Behind these processTranslateInst() methods, we'll still use the old transcriber logic to generate the IR, but the type & function lookups should be much simpler thanks to annotations.

Utility classes

• These classes should be instantiated inside the translator class as necessary. The main pass is not responsible for this.
• DifferentiableTypeContext: Parses all IRDifferentiableTypeAnnotation instructions in the current scope (involves walking up the parent tree, and reading all type annotations seen so far).
  • lookupWitnessForType(IRType*) should return the IDifferentiable or IDifferentiablePtrType witness table for the given type by performing a lookup.
  • Miscellaneous methods on this witness (looking up dadd, dzero, etc..), as well as synthesizing an existential dadd if necessary.
  • For now, we'll keep the rest the same, though we do need to eventually change this up.
• AssociatedFuncContext: Parses all IRAssociatedFuncAnnotation instructions in the current scope (involves walking up the parent tree, and reading all function annotations seen so far).
  • getAssociatedInst(IRInst* inst, IRFunctionOperator*) should return the associated callee for the given instruction & operator.

Stuff that can be removed

• IRDifferentiableTypeDictionary (and IRDifferentiableTypeDictionaryEntry)
• All the decorators reponsible for recording differentiability & links between custom derivatives.
  (Custom-derivatives are now entirely a front-end concern, with no need for an IR-level ops)
• Any logic to search for custom derivatives and coerce them into the right form.
• Any logic to handle existential functions and types. Both of these should automatically be resolved in the annotations
• dmul: this is just a cleanup from before. This method is not used anywhere, so its safe to remove.

[tangent-vector]

This is a really interesting direction to take these operators.

Seeing just how much this proposed direction is treating functions like types (holding witnesses for requirements, allowing extensions, etc.), I want to ask whether there's more that could be done in that direction.

To get really concrete, a couple provocative questions are:

• Could we allow the existing extension construct to be applied to a function?

• Could these requirements that a higher-order operator places on the functions be codified into actual interfaces, which specific function declarations could then conform to?

My hope would be that we could make all of this new functionality a bit less "special" if we can frame it in terms of taking existing language concepts that apply to types (extensions and interfaces), and allowing them to apply to functions as well.