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>
This commit is contained in:
Karan Janthe
2025-08-23 21:56:56 +00:00
parent e1258e79d6
commit 375e14ef49
7 changed files with 343 additions and 14 deletions

View File

@@ -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);