small cleanups, make Enzyme config nicer

Author: ZuseZ4

compiler/rustc_builtin_macros/messages.ftl

@@ -2,6 +2,7 @@ builtin_macros_alloc_error_must_be_fn = alloc_error_handler must be a function
 builtin_macros_alloc_must_statics = allocators must be statics
 
 builtin_macros_autodiff = autodiff must be applied to function
+builtin_macros_autodiff_not_build = this rustc version does not support autodiff
 
 builtin_macros_asm_clobber_abi = clobber_abi
 builtin_macros_asm_clobber_no_reg = asm with `clobber_abi` must specify explicit registers for outputs

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -1,6 +1,4 @@
 #![allow(unused_imports)]
-#![allow(unused_variables)]
-#![allow(unused_mut)]
 //use crate::util::check_builtin_macro_attribute;
 //use crate::util::check_autodiff;
 
@@ -20,12 +18,25 @@ use rustc_span::Symbol;
 use std::string::String;
 use thin_vec::{thin_vec, ThinVec};
 
+#[cfg(llvm_enzyme)]
 fn first_ident(x: &NestedMetaItem) -> rustc_span::symbol::Ident {
     let segments = &x.meta_item().unwrap().path.segments;
     assert!(segments.len() == 1);
     segments[0].ident
 }
 
+#[cfg(not(llvm_enzyme))]
+pub fn expand(
+    ecx: &mut ExtCtxt<'_>,
+    _expand_span: Span,
+    meta_item: &ast::MetaItem,
+    item: Annotatable,
+) -> Vec<Annotatable> {
+    ecx.sess.dcx().emit_err(errors::AutoDiffSupportNotBuild { span: meta_item.span });
+    return vec![item];
+}
+
+#[cfg(llvm_enzyme)]
 pub fn expand(
     ecx: &mut ExtCtxt<'_>,
     expand_span: Span,
@@ -45,24 +56,16 @@ pub fn expand(
     let primal = orig_item.ident.clone();
 
     // Allow using `#[autodiff(...)]` only on a Fn
-    let (fn_item, has_ret, sig, sig_span) = if let Annotatable::Item(item) = &item
+    let (has_ret, sig, sig_span) = if let Annotatable::Item(item) = &item
         && let ItemKind::Fn(box ast::Fn { sig, .. }) = &item.kind
     {
-        (item, sig.decl.output.has_ret(), sig, ecx.with_call_site_ctxt(sig.span))
+        (sig.decl.output.has_ret(), sig, ecx.with_call_site_ctxt(sig.span))
     } else {
         ecx.sess.dcx().emit_err(errors::AutoDiffInvalidApplication { span: item.span() });
         return vec![item];
     };
     // create TokenStream from vec elemtents:
     // meta_item doesn't have a .tokens field
-    let ts: Vec<Token> = meta_item_vec.clone()[1..]
-        .iter()
-        .map(|x| {
-            let val = first_ident(x);
-            let t = Token::from_ast_ident(val);
-            t
-        })
-        .collect();
     let comma: Token = Token::new(TokenKind::Comma, Span::default());
     let mut ts: Vec<TokenTree> = vec![];
     for t in meta_item_vec.clone()[1..].iter() {
@@ -77,18 +80,15 @@ pub fn expand(
     dbg!(&x);
     let span = ecx.with_def_site_ctxt(expand_span);
 
-    let (d_sig, old_names, new_args, idents) = gen_enzyme_decl(&sig, &x, span);
+    let (d_sig, new_args, idents) = gen_enzyme_decl(&sig, &x, span);
     let new_decl_span = d_sig.span;
     let d_body = gen_enzyme_body(
         ecx,
         primal,
-        &old_names,
         &new_args,
         span,
         sig_span,
         new_decl_span,
-        &sig,
-        &d_sig,
         idents,
     );
     let d_ident = meta_item_vec[0].meta_item().unwrap().path.segments[0].ident;
@@ -102,7 +102,7 @@ pub fn expand(
     }));
     let mut rustc_ad_attr =
         P(ast::NormalAttr::from_ident(Ident::with_dummy_span(sym::rustc_autodiff)));
-    let mut attr: ast::Attribute = ast::Attribute {
+    let attr: ast::Attribute = ast::Attribute {
         kind: ast::AttrKind::Normal(rustc_ad_attr.clone()),
         id: ast::AttrId::from_u32(0),
         style: ast::AttrStyle::Outer,
@@ -116,7 +116,7 @@ pub fn expand(
         delim: rustc_ast::token::Delimiter::Parenthesis,
         tokens: ts,
     });
-    let mut attr2: ast::Attribute = ast::Attribute {
+    let attr2: ast::Attribute = ast::Attribute {
         kind: ast::AttrKind::Normal(rustc_ad_attr),
         id: ast::AttrId::from_u32(0),
         style: ast::AttrStyle::Outer,
@@ -131,6 +131,7 @@ pub fn expand(
 }
 
 // shadow arguments must be mutable references or ptrs, because Enzyme will write into them.
+#[cfg(llvm_enzyme)]
 fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
     let mut ty = ty.clone();
     match ty.kind {
@@ -152,37 +153,21 @@ fn assure_mut_ref(ty: &ast::Ty) -> ast::Ty {
 // The second will just take a tuple containing the new arguments.
 // This way we surpress rustc from optimizing any argument away.
 // The last line will 'loop {}', to match the return type of the new function
+#[cfg(llvm_enzyme)]
 fn gen_enzyme_body(
     ecx: &ExtCtxt<'_>,
     primal: Ident,
-    old_names: &[String],
     new_names: &[String],
     span: Span,
     sig_span: Span,
     new_decl_span: Span,
-    sig: &ast::FnSig,
-    d_sig: &ast::FnSig,
     idents: Vec<Ident>,
 ) -> P<ast::Block> {
     let blackbox_path = ecx.std_path(&[Symbol::intern("hint"), Symbol::intern("black_box")]);
-    let zeroed_path = ecx.std_path(&[Symbol::intern("mem"), Symbol::intern("zeroed")]);
     let empty_loop_block = ecx.block(span, ThinVec::new());
     let loop_expr = ecx.expr_loop(span, empty_loop_block);
-
     let blackbox_call_expr = ecx.expr_path(ecx.path(span, blackbox_path));
-    let zeroed_call_expr = ecx.expr_path(ecx.path(span, zeroed_path));
-
-    let mem_zeroed_call: Stmt =
-        ecx.stmt_expr(ecx.expr_call(span, zeroed_call_expr.clone(), thin_vec![]));
-    let unsafe_block_with_zeroed_call: P<ast::Expr> = ecx.expr_block(P(ast::Block {
-        stmts: thin_vec![mem_zeroed_call],
-        id: ast::DUMMY_NODE_ID,
-        rules: ast::BlockCheckMode::Unsafe(ast::UserProvided),
-        span: sig_span,
-        tokens: None,
-        could_be_bare_literal: false,
-    }));
-    let primal_call = gen_primal_call(ecx, span, primal, sig, idents);
+    let primal_call = gen_primal_call(ecx, span, primal, idents);
     // create ::core::hint::black_box(array(arr));
     let black_box_primal_call =
         ecx.expr_call(new_decl_span, blackbox_call_expr.clone(), thin_vec![primal_call.clone()]);
@@ -207,11 +192,11 @@ fn gen_enzyme_body(
     body
 }
 
+#[cfg(llvm_enzyme)]
 fn gen_primal_call(
     ecx: &ExtCtxt<'_>,
     span: Span,
     primal: Ident,
-    sig: &ast::FnSig,
     idents: Vec<Ident>,
 ) -> P<ast::Expr> {
     let primal_call_expr = ecx.expr_path(ecx.path_ident(span, primal));
@@ -226,17 +211,18 @@ fn gen_primal_call(
 // zero-initialized by Enzyme). Active arguments are not handled yet.
 // Each argument of the primal function (and the return type if existing) must be annotated with an
 // activity.
+#[cfg(llvm_enzyme)]
 fn gen_enzyme_decl(
     sig: &ast::FnSig,
     x: &AutoDiffAttrs,
     span: Span,
-) -> (ast::FnSig, Vec<String>, Vec<String>, Vec<Ident>) {
+) -> (ast::FnSig, Vec<String>, Vec<Ident>) {
     assert!(sig.decl.inputs.len() == x.input_activity.len());
     assert!(sig.decl.output.has_ret() == x.has_ret_activity());
     let mut d_decl = sig.decl.clone();
     let mut d_inputs = Vec::new();
     let mut new_inputs = Vec::new();
-    let mut old_names = Vec::new();
+    //let mut old_names = Vec::new();
     let mut idents = Vec::new();
     let mut act_ret = ThinVec::new();
     for (arg, activity) in sig.decl.inputs.iter().zip(x.input_activity.iter()) {
@@ -256,7 +242,7 @@ fn gen_enzyme_decl(
                     dbg!(&shadow_arg.pat);
                     panic!("not an ident?");
                 };
-                old_names.push(old_name.to_string());
+                //old_names.push(old_name.to_string());
                 let name: String = match x.mode {
                     DiffMode::Reverse => format!("d{}", old_name),
                     DiffMode::Forward => format!("b{}", old_name),
@@ -320,7 +306,7 @@ fn gen_enzyme_decl(
     // return type. This might require changing the return type to a
     // tuple.
     if act_ret.len() > 0 {
-        let mut ret_ty = match d_decl.output {
+        let ret_ty = match d_decl.output {
             FnRetTy::Ty(ref ty) => {
                 act_ret.insert(0, ty.clone());
                 let kind = TyKind::Tup(act_ret);
@@ -339,5 +325,5 @@ fn gen_enzyme_decl(
     }
 
     let d_sig = FnSig { header: sig.header.clone(), decl: d_decl, span };
-    (d_sig, old_names, new_inputs, idents)
+    (d_sig, new_inputs, idents)
 }

compiler/rustc_builtin_macros/src/errors.rs

@@ -171,6 +171,13 @@ pub(crate) struct AutoDiffInvalidApplication {
     pub(crate) span: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag(builtin_macros_autodiff_not_build)]
+pub(crate) struct AutoDiffSupportNotBuild {
+    #[primary_span]
+    pub(crate) span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(builtin_macros_concat_bytes_invalid)]
 pub(crate) struct ConcatBytesInvalid {

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1,6 +1,6 @@
 #![allow(non_camel_case_types)]
 #![allow(non_upper_case_globals)]
-#![allow(unexpected_cfgs)]
+//#![allow(unexpected_cfgs)]
 
 use rustc_ast::expand::autodiff_attrs::DiffActivity;
 
@@ -2609,10 +2609,10 @@ extern "C" {
     ) -> *mut c_void;
 }
 
-#[cfg(autodiff_fallback)]
+#[cfg(not(llvm_enzyme))]
 pub use self::Fallback_AD::*;
 
-#[cfg(autodiff_fallback)]
+#[cfg(not(llvm_enzyme))]
 pub mod Fallback_AD {
     #![allow(unused_variables)]
     use super::*;
@@ -2745,9 +2745,9 @@ pub mod Shared_AD {
     use libc::size_t;
     use super::Context;
 
-    #[cfg(autodiff_fallback)]
+    #[cfg(not(llvm_enzyme))]
     use super::Fallback_AD::*;
-    #[cfg(not(autodiff_fallback))]
+    #[cfg(llvm_enzyme)]
     use super::Enzyme_AD::*;
 
     use core::fmt;
@@ -2931,13 +2931,13 @@ pub mod Shared_AD {
     }
 }
 
-#[cfg(not(autodiff_fallback))]
+#[cfg(llvm_enzyme)]
 pub use self::Enzyme_AD::*;
 
 // Enzyme is an optional component, so we do need to provide a fallback when it is ont getting
 // compiled. We deny the usage of #[autodiff(..)] on a higher level, so a placeholder implementation
 // here is completely fine.
-#[cfg(not(autodiff_fallback))]
+#[cfg(llvm_enzyme)]
 pub mod Enzyme_AD {
 use super::*;
 

compiler/rustc_session/src/config.rs

@@ -1273,6 +1273,7 @@ fn default_configuration(sess: &Session) -> Cfg {
     //
     // NOTE: These insertions should be kept in sync with
     // `CheckCfg::fill_well_known` below.
+    ins_none!(sym::autodiff_fallback);
 
     if sess.opts.debug_assertions {
         ins_none!(sym::debug_assertions);
@@ -1460,6 +1461,7 @@ impl CheckCfg {
         //
         // When adding a new config here you should also update
         // `tests/ui/check-cfg/well-known-values.rs`.
+        ins!(sym::autodiff_fallback, no_values);
 
         ins!(sym::debug_assertions, no_values);
 

compiler/rustc_span/src/symbol.rs

@@ -440,6 +440,7 @@ symbols! {
         augmented_assignments,
         auto_traits,
         autodiff,
+        autodiff_fallback,
         automatically_derived,
         avx,
         avx512_target_feature,
@@ -493,6 +494,7 @@ symbols! {
         cfg_accessible,
         cfg_attr,
         cfg_attr_multi,
+        cfg_autodiff_fallback,
         cfg_doctest,
         cfg_eval,
         cfg_hide,

src/bootstrap/src/core/build_steps/compile.rs

@@ -1048,6 +1048,9 @@ pub fn rustc_cargo_env(
     if builder.config.rust_verify_llvm_ir {
         cargo.env("RUSTC_VERIFY_LLVM_IR", "1");
     }
+    if builder.config.llvm_enzyme {
+        cargo.rustflag("--cfg=llvm_enzyme");
+    }
 
     // Note that this is disabled if LLVM itself is disabled or we're in a check
     // build. If we are in a check build we still go ahead here presuming we've

src/bootstrap/src/core/config/config.rs

@@ -1095,6 +1095,7 @@ define_config! {
         codegen_backends: Option<Vec<String>> = "codegen-backends",
         lld: Option<bool> = "lld",
         lld_mode: Option<LldMode> = "use-lld",
+        llvm_enzyme: Option<bool> = "llvm-enzyme",
         llvm_tools: Option<bool> = "llvm-tools",
         deny_warnings: Option<bool> = "deny-warnings",
         backtrace_on_ice: Option<bool> = "backtrace-on-ice",
@@ -1545,6 +1546,7 @@ impl Config {
                 save_toolstates,
                 codegen_backends,
                 lld,
+                llvm_enzyme,
                 llvm_tools,
                 deny_warnings,
                 backtrace_on_ice,
@@ -1634,6 +1636,8 @@ impl Config {
             }
 
             set(&mut config.llvm_tools_enabled, llvm_tools);
+            config.llvm_enzyme =
+                llvm_enzyme.unwrap_or(config.channel == "dev" || config.channel == "nightly");
             config.rustc_parallel =
                 parallel_compiler.unwrap_or(config.channel == "dev" || config.channel == "nightly");
             config.rustc_default_linker = default_linker;

src/bootstrap/src/lib.rs

@@ -76,6 +76,9 @@ const LLD_FILE_NAMES: &[&str] = &["ld.lld", "ld64.lld", "lld-link", "wasm-ld"];
 /// (Mode restriction, config name, config values (if any))
 const EXTRA_CHECK_CFGS: &[(Option<Mode>, &str, Option<&[&'static str]>)] = &[
     (None, "bootstrap", None),
+    (Some(Mode::Rustc), "llvm_enzyme", None),
+    (Some(Mode::Codegen), "llvm_enzyme", None),
+    (Some(Mode::ToolRustc), "llvm_enzyme", None),
     (Some(Mode::Rustc), "parallel_compiler", None),
     (Some(Mode::ToolRustc), "parallel_compiler", None),
     (Some(Mode::ToolRustc), "rust_analyzer", None),