diff options
| author | Karan Janthe <karanjanthe@gmail.com> | 2025-08-23 21:56:56 +0000 |
|---|---|---|
| committer | Karan Janthe <karanjanthe@gmail.com> | 2025-09-19 04:02:19 +0000 |
| commit | 375e14ef491ac7bfa701e269b2815625abf2fca6 (patch) | |
| tree | 398d7a2d0dbe09de910c73061a84c210d55455b4 /compiler/rustc_codegen_llvm/src/builder/autodiff.rs | |
| parent | e1258e79d6cb709b26ded97d32de6c55f355e2aa (diff) | |
| download | rust-375e14ef491ac7bfa701e269b2815625abf2fca6.tar.gz rust-375e14ef491ac7bfa701e269b2815625abf2fca6.zip | |
Add TypeTree metadata attachment for autodiff
- 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 <karanjanthe@gmail.com>
Diffstat (limited to 'compiler/rustc_codegen_llvm/src/builder/autodiff.rs')
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index b66e3dfdeec..c3485f56391 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -1,6 +1,7 @@ use std::ptr; use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode}; +use rustc_ast::expand::typetree::FncTree; use rustc_codegen_ssa::common::TypeKind; use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods}; use rustc_middle::ty::{Instance, PseudoCanonicalInput, TyCtxt, TypingEnv}; @@ -294,6 +295,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( fn_args: &[&'ll Value], attrs: AutoDiffAttrs, dest: PlaceRef<'tcx, &'ll Value>, + fnc_tree: FncTree, ) { // We have to pick the name depending on whether we want forward or reverse mode autodiff. let mut ad_name: String = match attrs.mode { @@ -370,6 +372,10 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( fn_args, ); + if !fnc_tree.args.is_empty() || !fnc_tree.ret.0.is_empty() { + crate::typetree::add_tt(cx.llmod, cx.llcx, fn_to_diff, fnc_tree); + } + let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None); builder.store_to_place(call, dest.val); |
