We present Enzyme MLIR for Reverse Mode AD, which is designed to be easily modifiable to generate fast gradients for arbitrary dialects.
To our knowledge there exists no AD framework in MLIR which is easily extendable to new dialects currently.
AD in MLIR can unite the benefits of AD on LLVM-IR and the AST, where it is possible to make use of compiler optimizations, but also perform optimizations which require higher-level information which is lost in the lowering to LLVM-IR such as matrix multiplications.
In order to be independent of specific dialects we place heavy emphasis on well designed interfaces, such that custom operations are emitted e.g. for storing values for the backward pass which can then be lowered to arbitrary dialects.
This also allows the independent development of AD specific optimizations, which can then be exposed in MLIR-frontends and be shared between different dialects.