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}