-
Notifications
You must be signed in to change notification settings - Fork 13.8k
TypeTree support in autodiff #144197
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
base: master
Are you sure you want to change the base?
TypeTree support in autodiff #144197
Conversation
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Currently, I have implemented only for |
r? @ZuseZ4 |
Some changes occurred in compiler/rustc_ast/src/expand/autodiff_attrs.rs cc @ZuseZ4 Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs cc @ZuseZ4 Some changes occurred in compiler/rustc_codegen_ssa Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs cc @ZuseZ4 |
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
This comment has been minimized.
Some changes occurred in compiler/rustc_codegen_gcc |
This comment has been minimized.
This comment has been minimized.
CI is failing, fixing them! |
Some changes occurred in src/tools/enzyme cc @ZuseZ4 |
This comment has been minimized.
This comment has been minimized.
d3f78d7
to
94aa8bc
Compare
This comment has been minimized.
This comment has been minimized.
@ZuseZ4 I have updated the tests and also fixed recursion, need to update the enzyme submodule in order to fix it, but you can review it |
This comment has been minimized.
This comment has been minimized.
8693753
to
6339a2c
Compare
☔ The latest upstream changes (presumably #146671) made this pull request unmergeable. Please resolve the merge conflicts. |
Signed-off-by: Karan Janthe <[email protected]>
- Add F128 support to TypeTree Kind enum - Implement TypeTree FFI bindings and conversion functions - Add typetree.rs module for metadata attachment to LLVM functions - Integrate TypeTree generation with autodiff intrinsic pipeline - Support scalar types: f32, f64, integers, f16, f128 - Attach enzyme_type attributes as LLVM string metadata for Enzyme Signed-off-by: Karan Janthe <[email protected]>
- Fix nott-flag test to emit LLVM IR and check enzyme_type attributes - Replace TODO comments with actual TypeTree metadata verification - Test that NoTT flag properly disables TypeTree generation - Test that TypeTree enabled generates proper enzyme_type attributes Signed-off-by: Karan Janthe <[email protected]>
- Add specific tests for f32, f64, i32, f16, f128 TypeTree generation - Verify correct enzyme_type metadata for each scalar type - Ensure TypeTree metadata matches expected Enzyme format Signed-off-by: Karan Janthe <[email protected]>
6339a2c
to
7470015
Compare
This PR was rebased onto a different master commit. Here's a range-diff highlighting what actually changed. Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers. |
This comment has been minimized.
This comment has been minimized.
cc0b6c4
to
2f742ea
Compare
This comment has been minimized.
This comment has been minimized.
2f742ea
to
0d04346
Compare
This comment has been minimized.
This comment has been minimized.
0d04346
to
7070455
Compare
This comment has been minimized.
This comment has been minimized.
Signed-off-by: Karan Janthe <[email protected]>
7070455
to
3ba5f19
Compare
TypeTrees for Autodiff
What are TypeTrees?
Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.
Structure
Example:
fn compute(x: &f32, data: &[f32]) -> f32
Input 0:
x: &f32
Input 1:
data: &[f32]
Output:
f32
Why Needed?
What Enzyme Does With This Information:
Without TypeTrees (current state):
With TypeTrees (our implementation):
TypeTrees - Offset and -1 Explained
Type Structure
Offset Values
Regular Offset (0, 4, 8, etc.)
Specific byte position within a structure
TypeTree for
&Point
(internal representation):Generates LLVM:
Offset -1 (Special: "Everywhere")
Means "this pattern repeats for ALL elements"
Example 1: Array
[f32; 100]
Instead of listing 100 separate Types with offsets
0,4,8,12...396
Example 2: Slice
&[i32]
Example 3: Mixed Structure