rustc_monomorphize/collector/
autodiff.rs

1use rustc_middle::bug;
2use rustc_middle::ty::{self, GenericArg, IntrinsicDef, TyCtxt};
3
4use crate::collector::{MonoItems, create_fn_mono_item};
5
6// Here, we force both primal and diff function to be collected in
7// mono so this does not interfere in `autodiff` intrinsics
8// codegen process. If they are unused, LLVM will remove them when
9// compiling with O3.
10pub(crate) fn collect_autodiff_fn<'tcx>(
11    tcx: TyCtxt<'tcx>,
12    instance: ty::Instance<'tcx>,
13    intrinsic: IntrinsicDef,
14    output: &mut MonoItems<'tcx>,
15) {
16    if intrinsic.name != rustc_span::sym::autodiff {
17        return;
18    };
19
20    collect_autodiff_fn_from_arg(instance.args[0], tcx, output);
21}
22
23fn collect_autodiff_fn_from_arg<'tcx>(
24    arg: GenericArg<'tcx>,
25    tcx: TyCtxt<'tcx>,
26    output: &mut MonoItems<'tcx>,
27) {
28    let (instance, span) = match arg.kind() {
29        ty::GenericArgKind::Type(ty) => match ty.kind() {
30            ty::FnDef(def_id, substs) => {
31                let span = tcx.def_span(def_id);
32                let instance = ty::Instance::expect_resolve(
33                    tcx,
34                    ty::TypingEnv::non_body_analysis(tcx, def_id),
35                    *def_id,
36                    substs,
37                    span,
38                );
39
40                (instance, span)
41            }
42            _ => bug!("expected autodiff function"),
43        },
44        _ => bug!("expected type when matching autodiff arg"),
45    };
46
47    output.push(create_fn_mono_item(tcx, instance, span));
48}
rustc_monomorphize/collector/autodiff.rs

rustc_monomorphize/collector/
autodiff.rs
