-
Notifications
You must be signed in to change notification settings - Fork 232
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
[Auto-diff] Overhaul auto-diff type tracking + Overhaul dynamic dispatch for differentiable functions #5866
base: master
Are you sure you want to change the base?
[Auto-diff] Overhaul auto-diff type tracking + Overhaul dynamic dispatch for differentiable functions #5866
Conversation
…veenb25/slang into fix-assoc-type-autodiff-2
…uePack, MakeValuePack)`
{ | ||
SpecializationOptions specOptions; | ||
specOptions.lowerWitnessLookups = true; | ||
specializeModule(targetProgram, irModule, codeGenContext->getSink(), specOptions); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Performing lowerWitnessLookups can open up new opportunities for specializations, so we need to run that specialization-optimization loop again.
This is really calling for cleaning up the specialization pass to be just a peephole optimization pass.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I believe specializeModule()
already runs the witness lowering in a loop with the rest of the specialization pass?
@@ -827,17 +860,72 @@ InstPair ForwardDiffTranscriber::transcribeCall(IRBuilder* builder, IRCall* orig | |||
} | |||
} | |||
} | |||
|
|||
{ | |||
// --WORKAROUND-- |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Put all these in a function, and do early returns instead of nested if.
This is a large overhaul PR containing several fixes and overhauls:
Overhauled differentiable type lowering:
IRDifferentiableTypeDictionaryDecoration
is deprecated in favor ofIRDifferentiableTypeAnnotation
instructions that can appear anywhere in the code (module/generic/function scopes) depending on the scope of the type. This makes it much simpler to lower differentiable run-time types.IRDifferentiableTypeAnnotation(type, witness)
instructions are generated by registering a 'type lowering hook' callback during IR lowering that checks the type against the list of differentiable types & generates an annotation for that type.IRDifferentiableTypeAnnotation
is hoistable, so it is emitted at the same parent as the type itself.DifferentiableTypeConformanceContext
, used by the auto-diff passes, scans the function being differentiated, as well as all parent scopes, to build the full list of differentiable types in the context of the function.Overhaul differentiation of dynamic types
IInterface
is a pair of two interface types:IInterface
&IDifferentiable
, while the differential pairs of existential types & associated types are replaced with lookups of a synthesized differential pair type.the derivative of
doThing()
is represented using associated types that are synthesized to represent the pairs.In the above example,
IFoo
is automatically extended to add the new type requirements (in addition to the derivative method requirements, which is already an existing feature)After the interface is modified, all available concrete witness tables are extended with synthesized methods to meet these new requirements.
Note that this type & method synthesis happens on the IR side during auto-diff pair lowering for now.
During auto-diff, these types are still represented as
DifferentialPair<LookupWitness(table, assoc_type_key)>
though it is not a proper type.Ideally, the next step is to perform this on the AST side when creating the method requirements for the differentiable methods.
Additional fixes
There are several other minor fixes to bugs that appeared once the auto-diff step was moved to after the dynamic-dispatch lowering step.
UseGraph
: During the primal availability step, primal instructions that are used in differential contexts are stored into an intermediate variables & loaded from later. Many instructions (such asExtractExistentialType
) cannot be stores, and so is considered a 'passthrough' inst that should be cloned in rather than stored. The previous system used a set of such 'use paths' calledUseChain
s to facilitate this, but it fails in cases where there is a common node in these paths. This PR introducesUseGraph
to represent & clone a full graph of instruction between a primal inst and its differential uses.Fixes: ##5829