bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

working dupv and dupvonly for fwd mode

Browse source
This commit is contained in:
Manuel Drehwald 2025-04-05 03:10:19 -04:00
parent d4f880f8ce
commit a68ae0cbc1
5 changed files with 107 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

40
compiler/rustc_ast/src/expand/autodiff_attrs.rs
View file

@ -50,8 +50,16 @@ pub enum DiffActivity {
/// with it.
Dual,
/// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
/// with it. It expects the shadow argument to be `width` times larger than the original
/// input/output.
Dualv,
/// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
/// with it. Drop the code which updates the original input/output for maximum performance.
DualOnly,
/// Forward Mode, Compute derivatives for this input/output and *overwrite* the shadow argument
/// with it. Drop the code which updates the original input/output for maximum performance.
/// It expects the shadow argument to be `width` times larger than the original input/output.
DualvOnly,
/// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
Duplicated,
/// Reverse Mode, Compute derivatives for this &T or *T input and *add* it to the shadow argument.
@ -59,7 +67,15 @@ pub enum DiffActivity {
DuplicatedOnly,
/// All Integers must be Const, but these are used to mark the integer which represents the
/// length of a slice/vec. This is used for safety checks on slices.
FakeActivitySize,
/// The integer (if given) specifies the size of the slice element in bytes.
FakeActivitySize(Option<u32>),
}
impl DiffActivity {
pub fn is_dual_or_const(&self) -> bool {
use DiffActivity::*;
matches!(self, |Dual| DualOnly | Dualv | DualvOnly | Const)
}
}
/// We generate one of these structs for each `#[autodiff(...)]` attribute.
#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
@ -131,11 +147,7 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
match mode {
DiffMode::Error => false,
DiffMode::Source => false,
DiffMode::Forward => {
activity == DiffActivity::Dual
|| activity == DiffActivity::DualOnly
|| activity == DiffActivity::Const
}
DiffMode::Forward => activity.is_dual_or_const(),
DiffMode::Reverse => {
activity == DiffActivity::Const
|| activity == DiffActivity::Active
@ -153,10 +165,8 @@ pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
pub fn valid_ty_for_activity(ty: &P<Ty>, activity: DiffActivity) -> bool {
use DiffActivity::*;
// It's always allowed to mark something as Const, since we won't compute derivatives wrt. it.
if matches!(activity, Const) {
return true;
}
if matches!(activity, Dual | DualOnly) {
// Dual variants also support all types.
if activity.is_dual_or_const() {
return true;
}
// FIXME(ZuseZ4) We should make this more robust to also
@ -172,9 +182,7 @@ pub fn valid_input_activity(mode: DiffMode, activity: DiffActivity) -> bool {
return match mode {
DiffMode::Error => false,
DiffMode::Source => false,
DiffMode::Forward => {
matches!(activity, Dual | DualOnly | Const)
}
DiffMode::Forward => activity.is_dual_or_const(),
DiffMode::Reverse => {
matches!(activity, Active | ActiveOnly | Duplicated | DuplicatedOnly | Const)
}
@ -189,10 +197,12 @@ impl Display for DiffActivity {
DiffActivity::Active => write!(f, "Active"),
DiffActivity::ActiveOnly => write!(f, "ActiveOnly"),
DiffActivity::Dual => write!(f, "Dual"),
DiffActivity::Dualv => write!(f, "Dualv"),
DiffActivity::DualOnly => write!(f, "DualOnly"),
DiffActivity::DualvOnly => write!(f, "DualvOnly"),
DiffActivity::Duplicated => write!(f, "Duplicated"),
DiffActivity::DuplicatedOnly => write!(f, "DuplicatedOnly"),
DiffActivity::FakeActivitySize => write!(f, "FakeActivitySize"),
DiffActivity::FakeActivitySize(s) => write!(f, "FakeActivitySize({:?})", s),
}
}
}
@ -220,7 +230,9 @@ impl FromStr for DiffActivity {
"ActiveOnly" => Ok(DiffActivity::ActiveOnly),
"Const" => Ok(DiffActivity::Const),
"Dual" => Ok(DiffActivity::Dual),
"Dualv" => Ok(DiffActivity::Dualv),
"DualOnly" => Ok(DiffActivity::DualOnly),
"DualvOnly" => Ok(DiffActivity::DualvOnly),
"Duplicated" => Ok(DiffActivity::Duplicated),
"DuplicatedOnly" => Ok(DiffActivity::DuplicatedOnly),
_ => Err(()),