Devdocs #1684
Devdocs #1684
Conversation
DEVDOCS.md
Outdated
|
|
||
| ## Overview | ||
|
|
||
| Enzyme-JAX is a C++ project that integrates the Enzyme automatic differentiation tool with JAX, enabling automatic differentiation of external C++ code within JAX. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. |
There was a problem hiding this comment.
at this point its grown a bit more than just enzyme. It also contains polygeist's high performance raising, parallelization, and cross compilation workflow; as well as a bunch of tensor/linear algebra/communication optiumizations
There was a problem hiding this comment.
| Enzyme-JAX is a C++ project that integrates the Enzyme automatic differentiation tool with JAX, enabling automatic differentiation of external C++ code within JAX. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. | |
| Enzyme-JAX is a C++ project whose original aim was to integrate the Enzyme automatic differentiation tool [1] with JAX, enabling automatic differentiation of external C++ code within JAX. It has since expanded to incorporate Polygeist's [2] high performance raising, parallelization, cross compilation workflow, as well as numerous tensor, linear algerba, and communication optimizations. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. |
There was a problem hiding this comment.
Thanks! Maybe this should better be moved to the main README?
| #### 2. **Passes** (`src/enzyme_ad/jax/Passes/`) | ||
| MLIR passes implement transformations and optimizations. | ||
|
|
||
| - Tablegen definitions in `src/enzyme_ad/jax/Passes/Passes.td` |
There was a problem hiding this comment.
can we add a bit more detail here?
for example enzymehloopt.cpp contains (nearly) all the stablehlo/linear algebra level transformations
we can also make a subsection for the (polygeist) raising passes?
|
|
||
| ## Testing | ||
|
|
||
| Run tests with: |
There was a problem hiding this comment.
maybe also add a link to what a lit test is (e.g. syntax is command at top, file below, with checks, link to full spec), where tests go/etc?
DEVDOCS.md
Outdated
|
|
||
| ### Quick Build | ||
| ```bash | ||
| bazel build --repo_env=CC=clang-18 --color=yes --copt=-fbracket-depth=1024 --host_copt=-fbracket-depth=1024 -c dbg :enzymexlamlir-opt |
There was a problem hiding this comment.
in princpile the bracket depth shouldnt be required here [and is added via the .bazelrc]
| ``` | ||
| enzymexlamlir-opt --enzyme-hlo-generate-td="patterns=and_pad_pad" --transform-interpreter --enzyme-hlo-remove-transform -allow-unregistered-dialect input.mlir | ||
| ``` | ||
|
|
There was a problem hiding this comment.
maybe also a fourth section for derivative rules ? [cc @avik-pal ]
DEVDOCS.md
Outdated
|
|
||
| ## Overview | ||
|
|
||
| Enzyme-JAX is a C++ project that integrates the Enzyme automatic differentiation tool with JAX, enabling automatic differentiation of external C++ code within JAX. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. |
There was a problem hiding this comment.
| Enzyme-JAX is a C++ project that integrates the Enzyme automatic differentiation tool with JAX, enabling automatic differentiation of external C++ code within JAX. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. | |
| Enzyme-JAX is a C++ project whose original aim was to integrate the Enzyme automatic differentiation tool [1] with JAX, enabling automatic differentiation of external C++ code within JAX. It has since expanded to incorporate Polygeist's [2] high performance raising, parallelization, cross compilation workflow, as well as numerous tensor, linear algerba, and communication optimizations. The project uses LLVM's MLIR framework for intermediate representation and transformation of code. |
| 2. Inherit from `mlir::OpRewritePattern<OpType>` | ||
| 3. Implement `matchAndRewrite()` method | ||
| 4. Register in `EnzymeHLOOptPass::runOnOperation()` | ||
| 5. Register as Transform operation in `TransformOps.td` |
There was a problem hiding this comment.
We can also say to add to primitives.py which I often forget
There was a problem hiding this comment.
Thanks! Feel free to just push, otherwise I can do it later. I didn't even know about primitives.py 🙈
I wrote some development information in a single file. Please chime in and correct/add more information. This probably also needs to be refactored to a proper wiki layout before merging.