rustc_ast/expand/
autodiff_attrs.rs

1//! This crate handles the user facing autodiff macro. For each `#[autodiff(...)]` attribute,
2//! we create an [`AutoDiffItem`] which contains the source and target function names. The source
3//! is the function to which the autodiff attribute is applied, and the target is the function
4//! getting generated by us (with a name given by the user as the first autodiff arg).
5
6use std::fmt::{self, Display, Formatter};
7use std::str::FromStr;
8
9use crate::expand::{Decodable, Encodable, HashStable_Generic};
10use crate::ptr::P;
11use crate::{Ty, TyKind};
12
13/// Forward and Reverse Mode are well known names for automatic differentiation implementations.
14/// Enzyme does support both, but with different semantics, see DiffActivity. The First variants
15/// are a hack to support higher order derivatives. We need to compute first order derivatives
16/// before we compute second order derivatives, otherwise we would differentiate our placeholder
17/// functions. The proper solution is to recognize and resolve this DAG of autodiff invocations,
18/// as it's already done in the C++ and Julia frontend of Enzyme.
19///
20/// Documentation for using [reverse](https://enzyme.mit.edu/rust/rev.html) and
21/// [forward](https://enzyme.mit.edu/rust/fwd.html) mode is available online.
22#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
23pub enum DiffMode {
24    /// No autodiff is applied (used during error handling).
25    Error,
26    /// The primal function which we will differentiate.
27    Source,
28    /// The target function, to be created using forward mode AD.
29    Forward,
30    /// The target function, to be created using reverse mode AD.
31    Reverse,
32}
33
34/// Dual and Duplicated (and their Only variants) are getting lowered to the same Enzyme Activity.
35/// However, under forward mode we overwrite the previous shadow value, while for reverse mode
36/// we add to the previous shadow value. To not surprise users, we picked different names.
37/// Dual numbers is also a quite well known name for forward mode AD types.
38#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
39pub enum DiffActivity {
40    /// Implicit or Explicit () return type, so a special case of Const.
41    None,
42    /// Don't compute derivatives with respect to this input/output.
43    Const,
44    /// Reverse Mode, Compute derivatives for this scalar input/output.
45    Active,
46    /// Reverse Mode, Compute derivatives for this scalar output, but don't compute
47    /// the original return value.
48    ActiveOnly,
49    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
50    /// with it.
51    Dual,
52    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
53    /// with it. It expects the shadow argument to be `width` times larger than the original
54    /// input/output.
55    Dualv,
56    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
57    /// with it. Drop the code which updates the original input/output for maximum performance.
58    DualOnly,
59    /// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
60    /// with it. Drop the code which updates the original input/output for maximum performance.
61    /// It expects the shadow argument to be `width` times larger than the original input/output.
62    DualvOnly,
63    /// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
64    Duplicated,
65    /// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
66    /// Drop the code which updates the original input for maximum performance.
67    DuplicatedOnly,
68    /// All Integers must be Const, but these are used to mark the integer which represents the
69    /// length of a slice/vec. This is used for safety checks on slices.
70    /// The integer (if given) specifies the size of the slice element in bytes.
71    FakeActivitySize(Option<u32>),
72}
73
74impl DiffActivity {
75    pub fn is_dual_or_const(&self) -> bool {
76        use DiffActivity::*;
77        matches!(self, |Dual| DualOnly | Dualv | DualvOnly | Const)
78    }
79}
80/// We generate one of these structs for each `#[autodiff(...)]` attribute.
81#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
82pub struct AutoDiffItem {
83    /// The name of the function getting differentiated
84    pub source: String,
85    /// The name of the function being generated
86    pub target: String,
87    pub attrs: AutoDiffAttrs,
88}
89
90#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
91pub struct AutoDiffAttrs {
92    /// Conceptually either forward or reverse mode AD, as described in various autodiff papers and
93    /// e.g. in the [JAX
94    /// Documentation](https://jax.readthedocs.io/en/latest/_tutorials/advanced-autodiff.html#how-it-s-made-two-foundational-autodiff-functions).
95    pub mode: DiffMode,
96    /// A user-provided, batching width. If not given, we will default to 1 (no batching).
97    /// Calling a differentiated, non-batched function through a loop 100 times is equivalent to:
98    /// - Calling the function 50 times with a batch size of 2
99    /// - Calling the function 25 times with a batch size of 4,
100    /// etc. A batched function takes more (or longer) arguments, and might be able to benefit from
101    /// cache locality, better re-usal of primal values, and other optimizations.
102    /// We will (before LLVM's vectorizer runs) just generate most LLVM-IR instructions `width`
103    /// times, so this massively increases code size. As such, values like 1024 are unlikely to
104    /// work. We should consider limiting this to u8 or u16, but will leave it at u32 for
105    /// experiments for now and focus on documenting the implications of a large width.
106    pub width: u32,
107    pub ret_activity: DiffActivity,
108    pub input_activity: Vec<DiffActivity>,
109}
110
111impl AutoDiffAttrs {
112    pub fn has_primal_ret(&self) -> bool {
113        matches!(self.ret_activity, DiffActivity::Active | DiffActivity::Dual)
114    }
115}
116
117impl DiffMode {
118    pub fn is_rev(&self) -> bool {
119        matches!(self, DiffMode::Reverse)
120    }
121    pub fn is_fwd(&self) -> bool {
122        matches!(self, DiffMode::Forward)
123    }
124}
125
126impl Display for DiffMode {
127    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
128        match self {
129            DiffMode::Error => write!(f, "Error"),
130            DiffMode::Source => write!(f, "Source"),
131            DiffMode::Forward => write!(f, "Forward"),
132            DiffMode::Reverse => write!(f, "Reverse"),
133        }
134    }
135}
136
137/// Active(Only) is valid in reverse-mode AD for scalar float returns (f16/f32/...).
138/// Dual(Only) is valid in forward-mode AD for scalar float returns (f16/f32/...).
139/// Const is valid for all cases and means that we don't compute derivatives wrt. this output.
140/// That usually means we have a &mut or *mut T output and compute derivatives wrt. that arg,
141/// but this is too complex to verify here. Also it's just a logic error if users get this wrong.
142pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
143    if activity == DiffActivity::None {
144        // Only valid if primal returns (), but we can't check that here.
145        return true;
146    }
147    match mode {
148        DiffMode::Error => false,
149        DiffMode::Source => false,
150        DiffMode::Forward => activity.is_dual_or_const(),
151        DiffMode::Reverse => {
152            activity == DiffActivity::Const
153                || activity == DiffActivity::Active
154                || activity == DiffActivity::ActiveOnly
155        }
156    }
157}
158
159/// For indirections (ptr/ref) we can't use Active, since Active allocates a shadow value
160/// for the given argument, but we generally can't know the size of such a type.
161/// For scalar types (f16/f32/f64/f128) we can use Active and we can't use Duplicated,
162/// since Duplicated expects a mutable ref/ptr and we would thus end up with a shadow value
163/// who is an indirect type, which doesn't match the primal scalar type. We can't prevent
164/// users here from marking scalars as Duplicated, due to type aliases.
165pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
166    use DiffActivity::*;
167    // It's always allowed to mark something as Const, since we won't compute derivatives wrt. it.
168    // Dual variants also support all types.
169    if activity.is_dual_or_const() {
170        return true;
171    }
172    // FIXME(ZuseZ4) We should make this more robust to also
173    // handle type aliases. Once that is done, we can be more restrictive here.
174    if matches!(activity, Active | ActiveOnly) {
175        return true;
176    }
177    matches!(ty.kind, TyKind::Ptr(_) | TyKind::Ref(..))
178        && matches!(activity, Duplicated | DuplicatedOnly)
179}
180pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
181    use DiffActivity::*;
182    return match mode {
183        DiffMode::Error => false,
184        DiffMode::Source => false,
185        DiffMode::Forward => activity.is_dual_or_const(),
186        DiffMode::Reverse => {
187            matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
188        }
189    };
190}
191
192impl Display for DiffActivity {
193    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
194        match self {
195            DiffActivity::None => write!(f, "None"),
196            DiffActivity::Const => write!(f, "Const"),
197            DiffActivity::Active => write!(f, "Active"),
198            DiffActivity::ActiveOnly => write!(f, "ActiveOnly"),
199            DiffActivity::Dual => write!(f, "Dual"),
200            DiffActivity::Dualv => write!(f, "Dualv"),
201            DiffActivity::DualOnly => write!(f, "DualOnly"),
202            DiffActivity::DualvOnly => write!(f, "DualvOnly"),
203            DiffActivity::Duplicated => write!(f, "Duplicated"),
204            DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
205            DiffActivity::FakeActivitySize(s) => write!(f, "FakeActivitySize({:?})", s),
206        }
207    }
208}
209
210impl FromStr for DiffMode {
211    type Err = ();
212
213    fn from_str(s: &str) -> Result<DiffMode, ()> {
214        match s {
215            "Error" => Ok(DiffMode::Error),
216            "Source" => Ok(DiffMode::Source),
217            "Forward" => Ok(DiffMode::Forward),
218            "Reverse" => Ok(DiffMode::Reverse),
219            _ => Err(()),
220        }
221    }
222}
223impl FromStr for DiffActivity {
224    type Err = ();
225
226    fn from_str(s: &str) -> Result<DiffActivity, ()> {
227        match s {
228            "None" => Ok(DiffActivity::None),
229            "Active" => Ok(DiffActivity::Active),
230            "ActiveOnly" => Ok(DiffActivity::ActiveOnly),
231            "Const" => Ok(DiffActivity::Const),
232            "Dual" => Ok(DiffActivity::Dual),
233            "Dualv" => Ok(DiffActivity::Dualv),
234            "DualOnly" => Ok(DiffActivity::DualOnly),
235            "DualvOnly" => Ok(DiffActivity::DualvOnly),
236            "Duplicated" => Ok(DiffActivity::Duplicated),
237            "DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),
238            _ => Err(()),
239        }
240    }
241}
242
243impl AutoDiffAttrs {
244    pub fn has_ret_activity(&self) -> bool {
245        self.ret_activity != DiffActivity::None
246    }
247    pub fn has_active_only_ret(&self) -> bool {
248        self.ret_activity == DiffActivity::ActiveOnly
249    }
250
251    pub const fn error() -> Self {
252        AutoDiffAttrs {
253            mode: DiffMode::Error,
254            width: 0,
255            ret_activity: DiffActivity::None,
256            input_activity: Vec::new(),
257        }
258    }
259    pub fn source() -> Self {
260        AutoDiffAttrs {
261            mode: DiffMode::Source,
262            width: 0,
263            ret_activity: DiffActivity::None,
264            input_activity: Vec::new(),
265        }
266    }
267
268    pub fn is_active(&self) -> bool {
269        self.mode != DiffMode::Error
270    }
271
272    pub fn is_source(&self) -> bool {
273        self.mode == DiffMode::Source
274    }
275    pub fn apply_autodiff(&self) -> bool {
276        !matches!(self.mode, DiffMode::Error | DiffMode::Source)
277    }
278
279    pub fn into_item(self, source: String, target: String) -> AutoDiffItem {
280        AutoDiffItem { source, target, attrs: self }
281    }
282}
283
284impl fmt::Display for AutoDiffItem {
285    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
286        write!(f, "Differentiating {} -> {}", self.source, self.target)?;
287        write!(f, " with attributes: {:?}", self.attrs)
288    }
289}