rustc_codegen_llvm/
typetree.rs

1use std::ffi::{CString, c_char};
2
3use rustc_ast::expand::typetree::{FncTree, Kind, TypeTree as RustTypeTree};
4
5use crate::attributes;
6use crate::context::FullCx;
7use crate::llvm::{self, EnzymeWrapper, Value};
8
9fn to_enzyme_typetree(
10    rust_typetree: &RustTypeTree,
11    _data_layout: &str,
12    llcx: &llvm::Context,
13) -> (llvm::TypeTree, Vec<llvm::TypeTree>) {
14    let mut enzyme_tt = llvm::TypeTree::new();
15    let extra_ints = process_typetree_recursive(&mut enzyme_tt, &rust_typetree, &[], llcx);
16
17    let mut int_vec = ::alloc::vec::Vec::new()vec![];
18    for _ in 0..extra_ints {
19        let mut int_tt = llvm::TypeTree::new();
20        int_tt.insert(&[0], llvm::CConcreteType::DT_Integer, llcx);
21        int_vec.push(int_tt);
22    }
23
24    (enzyme_tt, int_vec)
25}
26
27fn process_typetree_recursive(
28    enzyme_tt: &mut llvm::TypeTree,
29    rust_typetree: &RustTypeTree,
30    parent_indices: &[i64],
31    llcx: &llvm::Context,
32) -> u32 {
33    let mut extra_ints = 0;
34    for rust_type in &rust_typetree.0 {
35        let concrete_type = match rust_type.kind {
36            Kind::Anything => llvm::CConcreteType::DT_Anything,
37            Kind::Integer => llvm::CConcreteType::DT_Integer,
38            Kind::Pointer => llvm::CConcreteType::DT_Pointer,
39            Kind::RustSlice => llvm::CConcreteType::DT_Pointer,
40            Kind::Half => llvm::CConcreteType::DT_Half,
41            Kind::Float => llvm::CConcreteType::DT_Float,
42            Kind::Double => llvm::CConcreteType::DT_Double,
43            Kind::F128 => llvm::CConcreteType::DT_FP128,
44            Kind::Unknown => llvm::CConcreteType::DT_Unknown,
45        };
46
47        let mut indices = parent_indices.to_vec();
48        if !parent_indices.is_empty() {
49            indices.push(rust_type.offset as i64);
50        } else if rust_type.offset == -1 {
51            indices.push(-1);
52        } else {
53            indices.push(rust_type.offset as i64);
54        }
55
56        enzyme_tt.insert(&indices, concrete_type, llcx);
57
58        if #[allow(non_exhaustive_omitted_patterns)] match rust_type.kind {
    Kind::RustSlice => true,
    _ => false,
}matches!(rust_type.kind, Kind::RustSlice) {
59            // We lower slices to `ptr,int`, so add the int here.
60            extra_ints += 1;
61        }
62
63        if #[allow(non_exhaustive_omitted_patterns)] match rust_type.kind {
    Kind::Pointer | Kind::RustSlice => true,
    _ => false,
}matches!(rust_type.kind, Kind::Pointer | Kind::RustSlice)
64            && !rust_type.child.0.is_empty()
65        {
66            process_typetree_recursive(enzyme_tt, &rust_type.child, &indices, llcx);
67        }
68    }
69    extra_ints
70}
71
72// Describes all the locations in which we know how to apply an Enzyme TypeTree.
73enum TTLocation {
74    Definition,
75    Callsite,
76}
77
78#[cfg_attr(not(feature = "llvm_enzyme"), allow(unused))]
79pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: FncTree) {
80    // TypeTree processing uses functions from Enzyme, which we might not have available if we did
81    // not build this compiler with `llvm_enzyme`. This feature is not strictly necessary, but
82    // skipping this function increases the chance that Enzyme fails to compile some code.
83    // FIXME(autodiff): In the future we should conditionally run this function even without the
84    // `llvm_enzyme` feature, in case that libEnzyme was provided via rustup.
85    #[cfg(not(feature = "llvm_enzyme"))]
86    return;
87
88    let tcx = cx.tcx;
89    if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) {
90        return;
91    }
92    if tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::NoTT) {
93        return;
94    }
95
96    let llmod = cx.llmod;
97    let llcx = cx.llcx;
98    let inputs = tt.args;
99    let ret_tt: RustTypeTree = tt.ret;
100
101    let llvm_data_layout: *const c_char = unsafe { llvm::LLVMGetDataLayoutStr(&*llmod) };
102    let llvm_data_layout =
103        std::str::from_utf8(unsafe { std::ffi::CStr::from_ptr(llvm_data_layout) }.to_bytes())
104            .expect("got a non-UTF8 data-layout from LLVM");
105
106    let attr_name = "enzyme_type";
107    let c_attr_name = CString::new(attr_name).unwrap();
108
109    let tt_location: TTLocation =
110        if llvm::LLVMRustIsCall(fn_def) { TTLocation::Callsite } else { TTLocation::Definition };
111
112    let mut offset = 0;
113    for (i, input) in inputs.iter().enumerate() {
114        let (enzyme_tt, extra_ints) = to_enzyme_typetree(&input, llvm_data_layout, llcx);
115
116        // This scope is just a visual reminder that we *must* drop the enzyme_wrapper before
117        // we drop any typetrees (mainly enzyme_tt and extra_ints). Drop calls can not accept
118        // arguments like an enzyme_wrapper, so the typetree drop impl has to call get_instance
119        // on the static enzyme instance, which is behind a Mutex. Therefore we'd deadlock if we
120        // hold the enzyme_wrapper while dropping the typetrees.
121        {
122            let enzyme_wrapper = EnzymeWrapper::get_instance();
123            let c_str = enzyme_wrapper.tree_to_cstr(enzyme_tt.inner);
124
125            let attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
126            let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset);
127            // FIXME(autodiff): We currently know that this is correct for all the cases in which we
128            // call this function. But we should make it more robust for the future.
129            match tt_location {
130                TTLocation::Definition => {
131                    attributes::apply_to_llfn(fn_def, arg_pos, &[attr]);
132                }
133                TTLocation::Callsite => {
134                    attributes::apply_to_callsite(fn_def, arg_pos, &[attr]);
135                }
136            }
137            enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
138            for v in &extra_ints {
139                offset += 1;
140                let c_str = enzyme_wrapper.tree_to_cstr(v.inner);
141                let int_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
142                let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset);
143                match tt_location {
144                    TTLocation::Definition => {
145                        attributes::apply_to_llfn(fn_def, arg_pos, &[int_attr]);
146                    }
147                    TTLocation::Callsite => {
148                        attributes::apply_to_callsite(fn_def, arg_pos, &[int_attr]);
149                    }
150                }
151                enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
152            }
153        }
154    }
155    // We will only fail this if Rust types got lowered to LLVM in a way that we didn't predict.
156    // Error, so we can learn from our mistakes.
157    if #[allow(non_exhaustive_omitted_patterns)] match tt_location {
    TTLocation::Definition => true,
    _ => false,
}matches!(tt_location, TTLocation::Definition) {
158        let expected = offset as usize + inputs.len();
159        let actual = llvm::count_params(fn_def) as usize;
160        if expected != actual {
161            tcx.dcx().warn(::alloc::__export::must_use({
        ::alloc::fmt::format(format_args!("autodiff type-tree failure. We expected {0} LLVM argument(s), but the generated LLVM function has {1} parameter(s)",
                expected, actual))
    })format!(
162                "autodiff type-tree failure. We expected {expected} LLVM argument(s), \
163                 but the generated LLVM function has {actual} parameter(s)"
164            ));
165        }
166    }
167
168    // FIXME(autodiff): We should think more about what it means if a function returns a slice or
169    // other fat ptrs.
170    let (enzyme_tt, _extra_ints) = to_enzyme_typetree(&ret_tt, llvm_data_layout, llcx);
171    if ret_tt != RustTypeTree::new() {
172        let enzyme_wrapper = EnzymeWrapper::get_instance();
173        let c_str = enzyme_wrapper.tree_to_cstr(enzyme_tt.inner);
174        let ret_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str);
175        let arg_pos = llvm::AttributePlace::ReturnValue;
176        match tt_location {
177            TTLocation::Definition => {
178                attributes::apply_to_llfn(fn_def, arg_pos, &[ret_attr]);
179            }
180            TTLocation::Callsite => {
181                attributes::apply_to_callsite(fn_def, arg_pos, &[ret_attr]);
182            }
183        }
184        enzyme_wrapper.tree_to_string_free(c_str.as_ptr());
185    }
186}
rustc_codegen_llvm/typetree.rs

rustc_codegen_llvm/
typetree.rs
