add batching frontend

Author: ZuseZ4

compiler/rustc_ast/src/expand/batch_attrs.rs

@@ -0,0 +1,130 @@
+//! This crate handles the user facing autodiff macro. For each `#[autodiff(...)]` attribute,
+//! we create an [`BatchItem`] which contains the source and target function names. The source
+//! is the function to which the autodiff attribute is applied, and the target is the function
+//! getting generated by us (with a name given by the user as the first autodiff arg).
+
+use std::fmt::{self, Display, Formatter};
+use std::str::FromStr;
+
+use crate::expand::{Decodable, Encodable, HashStable_Generic};
+
+/// Documentation for using [reverse](https://enzyme.mit.edu/rust/rev.html) and
+/// [forward](https://enzyme.mit.edu/rust/fwd.html) mode is available online.
+#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
+pub enum BatchMode {
+    /// No vectorization is applied (used during error handling).
+    Error,
+    /// The primal function which we will vectorize.
+    Source,
+}
+
+#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
+pub enum BatchActivity {
+    /// Don't batch this argument
+    Const,
+    /// Unsafe.
+    Leaf,
+    /// Just receive this argument N times.
+    Vector,
+}
+/// We generate one of these structs for each `#[autodiff(...)]` attribute.
+#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
+pub struct BatchItem {
+    /// The name of the function getting differentiated
+    pub source: String,
+    /// The name of the function being generated
+    pub target: String,
+    pub attrs: BatchAttrs,
+}
+#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
+pub struct BatchAttrs {
+    /// Conceptually either forward or reverse mode AD, as described in various autodiff papers and
+    /// e.g. in the [JAX
+    /// Documentation](https://jax.readthedocs.io/en/latest/_tutorials/advanced-autodiff.html#how-it-s-made-two-foundational-autodiff-functions).
+    pub mode: BatchMode,
+    pub width: usize,
+    pub input_activity: Vec<BatchActivity>,
+}
+
+impl Display for BatchMode {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            BatchMode::Error => write!(f, "Error"),
+            BatchMode::Source => write!(f, "Source"),
+        }
+    }
+}
+
+impl Display for BatchActivity {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            BatchActivity::Const => write!(f, "Const"),
+            BatchActivity::Leaf => write!(f, "Leaf"),
+            BatchActivity::Vector => write!(f, "Vector"),
+        }
+    }
+}
+
+impl FromStr for BatchMode {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<BatchMode, ()> {
+        match s {
+            "Error" => Ok(BatchMode::Error),
+            "Source" => Ok(BatchMode::Source),
+            _ => Err(()),
+        }
+    }
+}
+impl FromStr for BatchActivity {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<BatchActivity, ()> {
+        match s {
+            "Const" => Ok(BatchActivity::Const),
+            "Leaf" => Ok(BatchActivity::Leaf),
+            "Vector" => Ok(BatchActivity::Vector),
+            _ => Err(()),
+        }
+    }
+}
+
+impl BatchAttrs {
+
+    pub fn error() -> Self {
+        BatchAttrs {
+            mode: BatchMode::Error,
+            width: 0,
+            input_activity: Vec::new(),
+        }
+    }
+    pub fn source() -> Self {
+        BatchAttrs {
+            mode: BatchMode::Source,
+            width: 0,
+            input_activity: Vec::new(),
+        }
+    }
+
+    pub fn is_active(&self) -> bool {
+        self.mode != BatchMode::Error
+    }
+
+    pub fn is_source(&self) -> bool {
+        self.mode == BatchMode::Source
+    }
+    pub fn apply_batch(&self) -> bool {
+        !matches!(self.mode, BatchMode::Error | BatchMode::Source)
+    }
+
+    pub fn into_item(self, source: String, target: String) -> BatchItem {
+        BatchItem { source, target, attrs: self }
+    }
+}
+
+impl fmt::Display for BatchItem {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Batching {} -> {}", self.source, self.target)?;
+        write!(f, " with attributes: {:?}", self.attrs)
+    }
+}

compiler/rustc_ast/src/expand/mod.rs

@@ -8,6 +8,7 @@ use crate::MetaItem;
 
 pub mod allocator;
 pub mod autodiff_attrs;
+pub mod batch_attrs;
 pub mod typetree;
 
 #[derive(Debug, Clone, Encodable, Decodable, HashStable_Generic)]

compiler/rustc_builtin_macros/messages.ftl

@@ -69,15 +69,24 @@ builtin_macros_assert_requires_boolean = macro requires a boolean expression as
 builtin_macros_assert_requires_expression = macro requires an expression as an argument
     .suggestion = try removing semicolon
 
-builtin_macros_autodiff = autodiff must be applied to function
+builtin_macros_batch = batch must be applied to a function
+builtin_macros_batch_missing_config = batch requires at least a name and mode
+builtin_macros_batch_mode = unknown Mode: `{$mode}`. Use `Forward` or `Reverse`
+builtin_macros_batch_mode_activity = {$act} can not be used in {$mode} Mode
+builtin_macros_batch_not_build = this rustc version does not support batch
+builtin_macros_batch_number_activities = expected {$expected} activities, but found {$found}
+builtin_macros_batch_ty_activity = {$act} can not be used for this type
+builtin_macros_batch_unknown_activity = did not recognize Activity: `{$act}`
+
+builtin_macros_autodiff = autodiff must be applied to a function
 builtin_macros_autodiff_missing_config = autodiff requires at least a name and mode
 builtin_macros_autodiff_mode = unknown Mode: `{$mode}`. Use `Forward` or `Reverse`
 builtin_macros_autodiff_mode_activity = {$act} can not be used in {$mode} Mode
 builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 builtin_macros_autodiff_number_activities = expected {$expected} activities, but found {$found}
 builtin_macros_autodiff_ty_activity = {$act} can not be used for this type
-
 builtin_macros_autodiff_unknown_activity = did not recognize Activity: `{$act}`
+
 builtin_macros_bad_derive_target = `derive` may only be applied to `struct`s, `enum`s and `union`s
     .label = not applicable here
     .label2 = not a `struct`, `enum` or `union`