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