rustc_codegen_llvm/builder/
autodiff.rs

1use std::ptr;
2
3use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
4use rustc_codegen_ssa::ModuleCodegen;
5use rustc_codegen_ssa::back::write::ModuleConfig;
6use rustc_codegen_ssa::common::TypeKind;
7use rustc_codegen_ssa::traits::BaseTypeCodegenMethods;
8use rustc_errors::FatalError;
9use rustc_middle::bug;
10use tracing::{debug, trace};
11
12use crate::back::write::llvm_err;
13use crate::builder::SBuilder;
14use crate::context::SimpleCx;
15use crate::declare::declare_simple_fn;
16use crate::errors::{AutoDiffWithoutEnable, LlvmError};
17use crate::llvm::AttributePlace::Function;
18use crate::llvm::{Metadata, True};
19use crate::value::Value;
20use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
21
22fn get_params(fnc: &Value) -> Vec<&Value> {
23    let param_num = llvm::LLVMCountParams(fnc) as usize;
24    let mut fnc_args: Vec<&Value> = vec![];
25    fnc_args.reserve(param_num);
26    unsafe {
27        llvm::LLVMGetParams(fnc, fnc_args.as_mut_ptr());
28        fnc_args.set_len(param_num);
29    }
30    fnc_args
31}
32
33fn has_sret(fnc: &Value) -> bool {
34    let num_args = llvm::LLVMCountParams(fnc) as usize;
35    if num_args == 0 {
36        false
37    } else {
38        unsafe { llvm::LLVMRustHasAttributeAtIndex(fnc, 0, llvm::AttributeKind::StructRet) }
39    }
40}
41
42// When we call the `__enzyme_autodiff` or `__enzyme_fwddiff` function, we need to pass all the
43// original inputs, as well as metadata and the additional shadow arguments.
44// This function matches the arguments from the outer function to the inner enzyme call.
45//
46// This function also considers that Rust level arguments not always match the llvm-ir level
47// arguments. A slice, `&[f32]`, for example, is represented as a pointer and a length on
48// llvm-ir level. The number of activities matches the number of Rust level arguments, so we
49// need to match those.
50// FIXME(ZuseZ4): This logic is a bit more complicated than it should be, can we simplify it
51// using iterators and peek()?
52fn match_args_from_caller_to_enzyme<'ll>(
53    cx: &SimpleCx<'ll>,
54    width: u32,
55    args: &mut Vec<&'ll llvm::Value>,
56    inputs: &[DiffActivity],
57    outer_args: &[&'ll llvm::Value],
58    has_sret: bool,
59) {
60    debug!("matching autodiff arguments");
61    // We now handle the issue that Rust level arguments not always match the llvm-ir level
62    // arguments. A slice, `&[f32]`, for example, is represented as a pointer and a length on
63    // llvm-ir level. The number of activities matches the number of Rust level arguments, so we
64    // need to match those.
65    // FIXME(ZuseZ4): This logic is a bit more complicated than it should be, can we simplify it
66    // using iterators and peek()?
67    let mut outer_pos: usize = 0;
68    let mut activity_pos = 0;
69
70    if has_sret {
71        // Then the first outer arg is the sret pointer. Enzyme doesn't know about sret, so the
72        // inner function will still return something. We increase our outer_pos by one,
73        // and once we're done with all other args we will take the return of the inner call and
74        // update the sret pointer with it
75        outer_pos = 1;
76    }
77
78    let enzyme_const = cx.create_metadata("enzyme_const".to_string()).unwrap();
79    let enzyme_out = cx.create_metadata("enzyme_out".to_string()).unwrap();
80    let enzyme_dup = cx.create_metadata("enzyme_dup".to_string()).unwrap();
81    let enzyme_dupnoneed = cx.create_metadata("enzyme_dupnoneed".to_string()).unwrap();
82
83    while activity_pos < inputs.len() {
84        let diff_activity = inputs[activity_pos as usize];
85        // Duplicated arguments received a shadow argument, into which enzyme will write the
86        // gradient.
87        let (activity, duplicated): (&Metadata, bool) = match diff_activity {
88            DiffActivity::None => panic!("not a valid input activity"),
89            DiffActivity::Const => (enzyme_const, false),
90            DiffActivity::Active => (enzyme_out, false),
91            DiffActivity::ActiveOnly => (enzyme_out, false),
92            DiffActivity::Dual => (enzyme_dup, true),
93            DiffActivity::DualOnly => (enzyme_dupnoneed, true),
94            DiffActivity::Duplicated => (enzyme_dup, true),
95            DiffActivity::DuplicatedOnly => (enzyme_dupnoneed, true),
96            DiffActivity::FakeActivitySize => (enzyme_const, false),
97        };
98        let outer_arg = outer_args[outer_pos];
99        args.push(cx.get_metadata_value(activity));
100        args.push(outer_arg);
101        if duplicated {
102            // We know that duplicated args by construction have a following argument,
103            // so this can not be out of bounds.
104            let next_outer_arg = outer_args[outer_pos + 1];
105            let next_outer_ty = cx.val_ty(next_outer_arg);
106            // FIXME(ZuseZ4): We should add support for Vec here too, but it's less urgent since
107            // vectors behind references (&Vec<T>) are already supported. Users can not pass a
108            // Vec by value for reverse mode, so this would only help forward mode autodiff.
109            let slice = {
110                if activity_pos + 1 >= inputs.len() {
111                    // If there is no arg following our ptr, it also can't be a slice,
112                    // since that would lead to a ptr, int pair.
113                    false
114                } else {
115                    let next_activity = inputs[activity_pos + 1];
116                    // We analyze the MIR types and add this dummy activity if we visit a slice.
117                    next_activity == DiffActivity::FakeActivitySize
118                }
119            };
120            if slice {
121                // A duplicated slice will have the following two outer_fn arguments:
122                // (..., ptr1, int1, ptr2, int2, ...). We add the following llvm-ir to our __enzyme call:
123                // (..., metadata! enzyme_dup, ptr, ptr, int1, ...).
124                // FIXME(ZuseZ4): We will upstream a safety check later which asserts that
125                // int2 >= int1, which means the shadow vector is large enough to store the gradient.
126                assert_eq!(cx.type_kind(next_outer_ty), TypeKind::Integer);
127
128                for i in 0..(width as usize) {
129                    let next_outer_arg2 = outer_args[outer_pos + 2 * (i + 1)];
130                    let next_outer_ty2 = cx.val_ty(next_outer_arg2);
131                    assert_eq!(cx.type_kind(next_outer_ty2), TypeKind::Pointer);
132                    let next_outer_arg3 = outer_args[outer_pos + 2 * (i + 1) + 1];
133                    let next_outer_ty3 = cx.val_ty(next_outer_arg3);
134                    assert_eq!(cx.type_kind(next_outer_ty3), TypeKind::Integer);
135                    args.push(next_outer_arg2);
136                }
137                args.push(cx.get_metadata_value(enzyme_const));
138                args.push(next_outer_arg);
139                outer_pos += 2 + 2 * width as usize;
140                activity_pos += 2;
141            } else {
142                // A duplicated pointer will have the following two outer_fn arguments:
143                // (..., ptr, ptr, ...). We add the following llvm-ir to our __enzyme call:
144                // (..., metadata! enzyme_dup, ptr, ptr, ...).
145                if matches!(diff_activity, DiffActivity::Duplicated | DiffActivity::DuplicatedOnly)
146                {
147                    assert_eq!(cx.type_kind(next_outer_ty), TypeKind::Pointer);
148                }
149                // In the case of Dual we don't have assumptions, e.g. f32 would be valid.
150                args.push(next_outer_arg);
151                outer_pos += 2;
152                activity_pos += 1;
153
154                // Now, if width > 1, we need to account for that
155                for _ in 1..width {
156                    let next_outer_arg = outer_args[outer_pos];
157                    args.push(next_outer_arg);
158                    outer_pos += 1;
159                }
160            }
161        } else {
162            // We do not differentiate with resprect to this argument.
163            // We already added the metadata and argument above, so just increase the counters.
164            outer_pos += 1;
165            activity_pos += 1;
166        }
167    }
168}
169
170// On LLVM-IR, we can luckily declare __enzyme_ functions without specifying the input
171// arguments. We do however need to declare them with their correct return type.
172// We already figured the correct return type out in our frontend, when generating the outer_fn,
173// so we can now just go ahead and use that. This is not always trivial, e.g. because sret.
174// Beyond sret, this article describes our challenges nicely:
175// <https://yorickpeterse.com/articles/the-mess-that-is-handling-structure-arguments-and-returns-in-llvm/>
176// I.e. (i32, f32) will get merged into i64, but we don't handle that yet.
177fn compute_enzyme_fn_ty<'ll>(
178    cx: &SimpleCx<'ll>,
179    attrs: &AutoDiffAttrs,
180    fn_to_diff: &'ll Value,
181    outer_fn: &'ll Value,
182) -> &'ll llvm::Type {
183    let fn_ty = cx.get_type_of_global(outer_fn);
184    let mut ret_ty = cx.get_return_type(fn_ty);
185
186    let has_sret = has_sret(outer_fn);
187
188    if has_sret {
189        // Now we don't just forward the return type, so we have to figure it out based on the
190        // primal return type, in combination with the autodiff settings.
191        let fn_ty = cx.get_type_of_global(fn_to_diff);
192        let inner_ret_ty = cx.get_return_type(fn_ty);
193
194        let void_ty = unsafe { llvm::LLVMVoidTypeInContext(cx.llcx) };
195        if inner_ret_ty == void_ty {
196            // This indicates that even the inner function has an sret.
197            // Right now I only look for an sret in the outer function.
198            // This *probably* needs some extra handling, but I never ran
199            // into such a case. So I'll wait for user reports to have a test case.
200            bug!("sret in inner function");
201        }
202
203        if attrs.width == 1 {
204            // Enzyme returns a struct of style:
205            // `{ original_ret(if requested), float, float, ... }`
206            let mut struct_elements = vec![];
207            if attrs.has_primal_ret() {
208                struct_elements.push(inner_ret_ty);
209            }
210            // Next, we push the list of active floats, since they will be lowered to `enzyme_out`,
211            // and therefore part of the return struct.
212            let param_tys = cx.func_params_types(fn_ty);
213            for (act, param_ty) in attrs.input_activity.iter().zip(param_tys) {
214                if matches!(act, DiffActivity::Active) {
215                    // Now find the float type at position i based on the fn_ty,
216                    // to know what (f16/f32/f64/...) to add to the struct.
217                    struct_elements.push(param_ty);
218                }
219            }
220            ret_ty = cx.type_struct(&struct_elements, false);
221        } else {
222            // First we check if we also have to deal with the primal return.
223            match attrs.mode {
224                DiffMode::Forward => match attrs.ret_activity {
225                    DiffActivity::Dual => {
226                        let arr_ty =
227                            unsafe { llvm::LLVMArrayType2(inner_ret_ty, attrs.width as u64 + 1) };
228                        ret_ty = arr_ty;
229                    }
230                    DiffActivity::DualOnly => {
231                        let arr_ty =
232                            unsafe { llvm::LLVMArrayType2(inner_ret_ty, attrs.width as u64) };
233                        ret_ty = arr_ty;
234                    }
235                    DiffActivity::Const => {
236                        todo!("Not sure, do we need to do something here?");
237                    }
238                    _ => {
239                        bug!("unreachable");
240                    }
241                },
242                DiffMode::Reverse => {
243                    todo!("Handle sret for reverse mode");
244                }
245                _ => {
246                    bug!("unreachable");
247                }
248            }
249        }
250    }
251
252    // LLVM can figure out the input types on it's own, so we take a shortcut here.
253    unsafe { llvm::LLVMFunctionType(ret_ty, ptr::null(), 0, True) }
254}
255
256/// When differentiating `fn_to_diff`, take a `outer_fn` and generate another
257/// function with expected naming and calling conventions[^1] which will be
258/// discovered by the enzyme LLVM pass and its body populated with the differentiated
259/// `fn_to_diff`. `outer_fn` is then modified to have a call to the generated
260/// function and handle the differences between the Rust calling convention and
261/// Enzyme.
262/// [^1]: <https://enzyme.mit.edu/getting_started/CallingConvention/>
263// FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to
264// cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
265fn generate_enzyme_call<'ll>(
266    cx: &SimpleCx<'ll>,
267    fn_to_diff: &'ll Value,
268    outer_fn: &'ll Value,
269    attrs: AutoDiffAttrs,
270) {
271    // We have to pick the name depending on whether we want forward or reverse mode autodiff.
272    let mut ad_name: String = match attrs.mode {
273        DiffMode::Forward => "__enzyme_fwddiff",
274        DiffMode::Reverse => "__enzyme_autodiff",
275        _ => panic!("logic bug in autodiff, unrecognized mode"),
276    }
277    .to_string();
278
279    // add outer_fn name to ad_name to make it unique, in case users apply autodiff to multiple
280    // functions. Unwrap will only panic, if LLVM gave us an invalid string.
281    let name = llvm::get_value_name(outer_fn);
282    let outer_fn_name = std::str::from_utf8(name).unwrap();
283    ad_name.push_str(outer_fn_name);
284
285    // Let us assume the user wrote the following function square:
286    //
287    // ```llvm
288    // define double @square(double %x) {
289    // entry:
290    //  %0 = fmul double %x, %x
291    //  ret double %0
292    // }
293    // ```
294    //
295    // The user now applies autodiff to the function square, in which case fn_to_diff will be `square`.
296    // Our macro generates the following placeholder code (slightly simplified):
297    //
298    // ```llvm
299    // define double @dsquare(double %x) {
300    //  ; placeholder code
301    //  return 0.0;
302    // }
303    // ```
304    //
305    // so our `outer_fn` will be `dsquare`. The unsafe code section below now removes the placeholder
306    // code and inserts an autodiff call. We also add a declaration for the __enzyme_autodiff call.
307    // Again, the arguments to all functions are slightly simplified.
308    // ```llvm
309    // declare double @__enzyme_autodiff_square(...)
310    //
311    // define double @dsquare(double %x) {
312    // entry:
313    //   %0 = tail call double (...) @__enzyme_autodiff_square(double (double)* nonnull @square, double %x)
314    //   ret double %0
315    // }
316    // ```
317    unsafe {
318        let enzyme_ty = compute_enzyme_fn_ty(cx, &attrs, fn_to_diff, outer_fn);
319
320        // FIXME(ZuseZ4): the CC/Addr/Vis values are best effort guesses, we should look at tests and
321        // think a bit more about what should go here.
322        let cc = llvm::LLVMGetFunctionCallConv(outer_fn);
323        let ad_fn = declare_simple_fn(
324            cx,
325            &ad_name,
326            llvm::CallConv::try_from(cc).expect("invalid callconv"),
327            llvm::UnnamedAddr::No,
328            llvm::Visibility::Default,
329            enzyme_ty,
330        );
331
332        // Otherwise LLVM might inline our temporary code before the enzyme pass has a chance to
333        // do it's work.
334        let attr = llvm::AttributeKind::NoInline.create_attr(cx.llcx);
335        attributes::apply_to_llfn(ad_fn, Function, &[attr]);
336
337        // first, remove all calls from fnc
338        let entry = llvm::LLVMGetFirstBasicBlock(outer_fn);
339        let br = llvm::LLVMRustGetTerminator(entry);
340        llvm::LLVMRustEraseInstFromParent(br);
341
342        let last_inst = llvm::LLVMRustGetLastInstruction(entry).unwrap();
343        let mut builder = SBuilder::build(cx, entry);
344
345        let num_args = llvm::LLVMCountParams(&fn_to_diff);
346        let mut args = Vec::with_capacity(num_args as usize + 1);
347        args.push(fn_to_diff);
348
349        let enzyme_primal_ret = cx.create_metadata("enzyme_primal_return".to_string()).unwrap();
350        if matches!(attrs.ret_activity, DiffActivity::Dual | DiffActivity::Active) {
351            args.push(cx.get_metadata_value(enzyme_primal_ret));
352        }
353        if attrs.width > 1 {
354            let enzyme_width = cx.create_metadata("enzyme_width".to_string()).unwrap();
355            args.push(cx.get_metadata_value(enzyme_width));
356            args.push(cx.get_const_i64(attrs.width as u64));
357        }
358
359        let has_sret = has_sret(outer_fn);
360        let outer_args: Vec<&llvm::Value> = get_params(outer_fn);
361        match_args_from_caller_to_enzyme(
362            &cx,
363            attrs.width,
364            &mut args,
365            &attrs.input_activity,
366            &outer_args,
367            has_sret,
368        );
369
370        let call = builder.call(enzyme_ty, ad_fn, &args, None);
371
372        // This part is a bit iffy. LLVM requires that a call to an inlineable function has some
373        // metadata attached to it, but we just created this code oota. Given that the
374        // differentiated function already has partly confusing metadata, and given that this
375        // affects nothing but the auttodiff IR, we take a shortcut and just steal metadata from the
376        // dummy code which we inserted at a higher level.
377        // FIXME(ZuseZ4): Work with Enzyme core devs to clarify what debug metadata issues we have,
378        // and how to best improve it for enzyme core and rust-enzyme.
379        let md_ty = cx.get_md_kind_id("dbg");
380        if llvm::LLVMRustHasMetadata(last_inst, md_ty) {
381            let md = llvm::LLVMRustDIGetInstMetadata(last_inst)
382                .expect("failed to get instruction metadata");
383            let md_todiff = cx.get_metadata_value(md);
384            llvm::LLVMSetMetadata(call, md_ty, md_todiff);
385        } else {
386            // We don't panic, since depending on whether we are in debug or release mode, we might
387            // have no debug info to copy, which would then be ok.
388            trace!("no dbg info");
389        }
390
391        // Now that we copied the metadata, get rid of dummy code.
392        llvm::LLVMRustEraseInstUntilInclusive(entry, last_inst);
393
394        if cx.val_ty(call) == cx.type_void() || has_sret {
395            if has_sret {
396                // This is what we already have in our outer_fn (shortened):
397                // define void @_foo(ptr <..> sret([32 x i8]) initializes((0, 32)) %0, <...>) {
398                //   %7 = call [4 x double] (...) @__enzyme_fwddiff_foo(ptr @square, metadata !"enzyme_width", i64 4, <...>)
399                //   <Here we are, we want to add the following two lines>
400                //   store [4 x double] %7, ptr %0, align 8
401                //   ret void
402                // }
403
404                // now store the result of the enzyme call into the sret pointer.
405                let sret_ptr = outer_args[0];
406                let call_ty = cx.val_ty(call);
407                if attrs.width == 1 {
408                    assert_eq!(cx.type_kind(call_ty), TypeKind::Struct);
409                } else {
410                    assert_eq!(cx.type_kind(call_ty), TypeKind::Array);
411                }
412                llvm::LLVMBuildStore(&builder.llbuilder, call, sret_ptr);
413            }
414            builder.ret_void();
415        } else {
416            builder.ret(call);
417        }
418
419        // Let's crash in case that we messed something up above and generated invalid IR.
420        llvm::LLVMRustVerifyFunction(
421            outer_fn,
422            llvm::LLVMRustVerifierFailureAction::LLVMAbortProcessAction,
423        );
424    }
425}
426
427pub(crate) fn differentiate<'ll>(
428    module: &'ll ModuleCodegen<ModuleLlvm>,
429    cgcx: &CodegenContext<LlvmCodegenBackend>,
430    diff_items: Vec<AutoDiffItem>,
431    _config: &ModuleConfig,
432) -> Result<(), FatalError> {
433    for item in &diff_items {
434        trace!("{}", item);
435    }
436
437    let diag_handler = cgcx.create_dcx();
438
439    let cx = SimpleCx::new(module.module_llvm.llmod(), module.module_llvm.llcx, cgcx.pointer_size);
440
441    // First of all, did the user try to use autodiff without using the -Zautodiff=Enable flag?
442    if !diff_items.is_empty()
443        && !cgcx.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable)
444    {
445        return Err(diag_handler.handle().emit_almost_fatal(AutoDiffWithoutEnable));
446    }
447
448    // Before dumping the module, we want all the TypeTrees to become part of the module.
449    for item in diff_items.iter() {
450        let name = item.source.clone();
451        let fn_def: Option<&llvm::Value> = cx.get_function(&name);
452        let Some(fn_def) = fn_def else {
453            return Err(llvm_err(
454                diag_handler.handle(),
455                LlvmError::PrepareAutoDiff {
456                    src: item.source.clone(),
457                    target: item.target.clone(),
458                    error: "could not find source function".to_owned(),
459                },
460            ));
461        };
462        debug!(?item.target);
463        let fn_target: Option<&llvm::Value> = cx.get_function(&item.target);
464        let Some(fn_target) = fn_target else {
465            return Err(llvm_err(
466                diag_handler.handle(),
467                LlvmError::PrepareAutoDiff {
468                    src: item.source.clone(),
469                    target: item.target.clone(),
470                    error: "could not find target function".to_owned(),
471                },
472            ));
473        };
474
475        generate_enzyme_call(&cx, fn_def, fn_target, item.attrs.clone());
476    }
477
478    // FIXME(ZuseZ4): support SanitizeHWAddress and prevent illegal/unsupported opts
479
480    trace!("done with differentiate()");
481
482    Ok(())
483}