Modifies activities for the AutoDiffOp. More...
Public Member Functions | |
| LogicalResult | matchAndRewrite (SourceOp uop, PatternRewriter &rewriter) const override |
Detailed Description
class ReverseRetOpt< SourceOp >
Modifies activities for the AutoDiffOp.
This is also a nice place to understand the semantics of the rev-mode autodiff op. At it's core, the reverse mode autodiff takes in a function
f: def f (pInput): ...perform computation return pOutput
One can assign a very simple function signature to f here: f: pInput -> pOutput
When trying to differentiate this function (using autodiff op), Enzyme creates a new function which also takes in the co-tangents of the outputs (dOutput), and computes and returns both the output and the input co-tangent (dInput). This is how the generated autodiff op eventually looks like:
def revdiff_f(pInput, dOutput): ...perform computation to compute pOutput ...perform computation to compute dInput return pOutput, dInput
The new function signature now becomes revdiff_f : (pInput', dOutput) -> (pOutput, dInput)
I mention pInput' here because it is not exactly the input arguments to the function we are differentiating, for example, if the input argument type is an enzyme_dup, we will provide both tht primal value along with the shadow(which we then accumulate into and return). So specifically,
pInput' = pInput (if the activity is enzyme_active, enzyme_const) | pInput, dInput (if the activity is enzyme_dup, enzyme_dupnoneed)
Now that we have fixed the codegen semantics, we can go ahead and optimize for both the input return activities based on usage. Possible activity promotion flow for the inputs can be as follows:
- enzyme_active --> enzyme_const (dInput is never used, so we simply don't compute it)
- enzyme_dup --> enzyme_const (pInput is mutable, readonly, nocapture, dInput is never used post AD)
Similarly, one can define a similar activity promotion flow for the outputs:
- enzyme_active --> enzyme_activenoneed (we do need to pass in dOutput into the function, but we can see that pOutput is never used, so let's just not return it. This has the advantage of triggering some additional DCE inside the generated derivative function)
- enzyme_const --> enzyme_constnoneed (same as above, but we now simply skip over this output)
One other thing to note here is that these optimizations preserve the input function signature, and only modify the number of outputs.
Member Function Documentation
◆ matchAndRewrite()
|
inlineoverride |
Definition at line 689 of file Ops.cpp.
References isMutable().
The documentation for this class was generated from the following file:
- MLIR/Dialect/Ops.cpp
Generated on Fri May 8 2026 19:56:26 for Enzyme by
1.12.0