From 087ffd73bf43671db22087570c97e0afdde20f31 Mon Sep 17 00:00:00 2001 From: Manuel Drehwald Date: Thu, 3 Apr 2025 17:19:11 -0400 Subject: add the autodiff batch mode frontend --- compiler/rustc_ast/src/expand/autodiff_attrs.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) (limited to 'compiler/rustc_ast') diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs index c8ec185ee5e..f01c781f46c 100644 --- a/compiler/rustc_ast/src/expand/autodiff_attrs.rs +++ b/compiler/rustc_ast/src/expand/autodiff_attrs.rs @@ -77,6 +77,17 @@ pub struct AutoDiffAttrs { /// 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: DiffMode, + /// A user-provided, batching width. If not given, we will default to 1 (no batching). + /// Calling a differentiated, non-batched function through a loop 100 times is equivalent to: + /// - Calling the function 50 times with a batch size of 2 + /// - Calling the function 25 times with a batch size of 4, + /// etc. A batched function takes more (or longer) arguments, and might be able to benefit from + /// cache locality, better re-usal of primal values, and other optimizations. + /// We will (before LLVM's vectorizer runs) just generate most LLVM-IR instructions `width` + /// times, so this massively increases code size. As such, values like 1024 are unlikely to + /// work. We should consider limiting this to u8 or u16, but will leave it at u32 for + /// experiments for now and focus on documenting the implications of a large width. + pub width: u32, pub ret_activity: DiffActivity, pub input_activity: Vec, } @@ -222,6 +233,7 @@ impl AutoDiffAttrs { pub const fn error() -> Self { AutoDiffAttrs { mode: DiffMode::Error, + width: 0, ret_activity: DiffActivity::None, input_activity: Vec::new(), } @@ -229,6 +241,7 @@ impl AutoDiffAttrs { pub fn source() -> Self { AutoDiffAttrs { mode: DiffMode::Source, + width: 0, ret_activity: DiffActivity::None, input_activity: Vec::new(), } -- cgit 1.4.1-3-g733a5