Skip to content

Implement autodiff using intrinsics #142640

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 44 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
44 commits
Select commit Hold shift + click to select a range
a68be0a
Lower autodiff functions using instrinsics
Sa4dUs Jun 4, 2025
bca8b29
Macro expansion with `rustc_intrinsic`
Sa4dUs Jun 5, 2025
c4857a0
Lowering draft
Sa4dUs Jun 17, 2025
3aff1f4
Naive impl of intrinsic codegen
Sa4dUs Jun 23, 2025
cabd40b
Feature intrinsics in cg tests
Sa4dUs Jun 24, 2025
b4ac54a
Remove `sret` logic
Sa4dUs Jun 29, 2025
6c484e9
Move logic to a dedicated `enzyme_autodiff` intrinsic
Sa4dUs Jul 7, 2025
880b852
Remove attr checking from hir_analysis
Sa4dUs Jul 7, 2025
cfa89c8
FIx generics error when passing fn as param to intrinsic
Sa4dUs Jul 8, 2025
78c7f01
Use Instance::new_raw instead of Instance::mono
Sa4dUs Jul 8, 2025
5a42b4a
Hacky fix for issues at trait calls
Sa4dUs Jul 10, 2025
60c9c0b
Fix how fns where being retrieved at intrinsic cg
Sa4dUs Jul 11, 2025
8e3ab9a
Use Self instead of Foo placeholder
Sa4dUs Jul 11, 2025
7f156e7
Remove unused code
Sa4dUs Jul 12, 2025
7b204bb
Remove primal call and collect it in mono instead
Sa4dUs Jul 14, 2025
072e54c
Update codegen tests
Sa4dUs Jul 15, 2025
31b3f21
Handle slices when extracting args from tuple
Sa4dUs Jul 22, 2025
21b40ca
Do not depend on mono anymore
Sa4dUs Jul 23, 2025
22518e0
Get args from tuple using fnabi and minor fixes
Sa4dUs Jul 23, 2025
34ca48a
Remove dead code
Sa4dUs Jul 25, 2025
791b791
Reintroduce autodiff enable and lto errors
Sa4dUs Jul 25, 2025
bff2dae
Minor fixes after rebase
Sa4dUs Jul 25, 2025
5a29edd
Add rest of test fixes
Sa4dUs Jul 25, 2025
572a9b7
Remove cfg enzyme for ci
Sa4dUs Jul 25, 2025
3aecee9
FIx cg_gcc
Sa4dUs Jul 25, 2025
0c22a2c
Macro expansion cleanup
Sa4dUs Jul 27, 2025
16e23b7
Remove dead code
Sa4dUs Jul 27, 2025
275aff9
Add `enzyme_autodiff` doc comment
Sa4dUs Jul 28, 2025
4e9fdeb
Add expansion example to intrinsic docs
Sa4dUs Jul 28, 2025
d6acc7d
Better error handling
Sa4dUs Jul 29, 2025
9fca70b
Allow `core_intrinsics` when `autodiff` is enabled
Sa4dUs Jul 29, 2025
9441c06
Remove autodiff limitations subsection
Sa4dUs Jul 30, 2025
cc2a43a
Remove inlining for autodiff handling
Sa4dUs Aug 7, 2025
ff7594a
Update tests for LLVM 21
Sa4dUs Aug 9, 2025
f72d7e3
Rename intrinsic from `enzyme_autodiff` to `autodiff`
Sa4dUs Aug 14, 2025
1009e39
Handle empty arg tuple
Sa4dUs Aug 14, 2025
e8919f4
Add turbofish comment
Sa4dUs Aug 14, 2025
0a8a3a7
Remove `gen_enzyme_body` and call it in place instead
Sa4dUs Aug 14, 2025
612f437
Remove lto inline logic
Sa4dUs Aug 14, 2025
198fbdf
Fix collector opt comment
Sa4dUs Aug 14, 2025
cd68456
Remove `normalize_erasing_late_bound_regions` and simply `.skip_binde…
Sa4dUs Aug 14, 2025
d553d57
Add macro hygiene test
Sa4dUs Aug 14, 2025
cca2d6d
Fix tidy
Sa4dUs Aug 14, 2025
fe9da82
Fix tidy 2
Sa4dUs Aug 14, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions Cargo.lock
Original file line number Diff line number Diff line change
Expand Up @@ -4307,7 +4307,6 @@ name = "rustc_monomorphize"
version = "0.0.0"
dependencies = [
"rustc_abi",
"rustc_ast",
"rustc_data_structures",
"rustc_errors",
"rustc_fluent_macro",
Expand All @@ -4316,7 +4315,6 @@ dependencies = [
"rustc_middle",
"rustc_session",
"rustc_span",
"rustc_symbol_mangling",
"rustc_target",
"serde",
"serde_json",
Expand Down
435 changes: 130 additions & 305 deletions compiler/rustc_builtin_macros/src/autodiff.rs

Large diffs are not rendered by default.

6 changes: 0 additions & 6 deletions compiler/rustc_codegen_gcc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,6 @@ use gccjit::{CType, Context, OptimizationLevel};
#[cfg(feature = "master")]
use gccjit::{TargetInfo, Version};
use rustc_ast::expand::allocator::AllocatorKind;
use rustc_ast::expand::autodiff_attrs::AutoDiffItem;
use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
use rustc_codegen_ssa::back::write::{
CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
Expand Down Expand Up @@ -363,12 +362,7 @@ impl WriteBackendMethods for GccCodegenBackend {
_exported_symbols_for_lto: &[String],
each_linked_rlib_for_lto: &[PathBuf],
modules: Vec<FatLtoInput<Self>>,
diff_functions: Vec<AutoDiffItem>,
) -> Result<ModuleCodegen<Self::Module>, FatalError> {
if !diff_functions.is_empty() {
unimplemented!();
}

back::lto::run_fat(cgcx, each_linked_rlib_for_lto, modules)
}

Expand Down
16 changes: 0 additions & 16 deletions compiler/rustc_codegen_llvm/src/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,22 +28,6 @@ pub(crate) fn apply_to_callsite(callsite: &Value, idx: AttributePlace, attrs: &[
}
}

pub(crate) fn has_attr(llfn: &Value, idx: AttributePlace, attr: AttributeKind) -> bool {
llvm::HasAttributeAtIndex(llfn, idx, attr)
}

pub(crate) fn has_string_attr(llfn: &Value, name: &str) -> bool {
llvm::HasStringAttribute(llfn, name)
}

pub(crate) fn remove_from_llfn(llfn: &Value, place: AttributePlace, kind: AttributeKind) {
llvm::RemoveRustEnumAttributeAtIndex(llfn, place, kind);
}

pub(crate) fn remove_string_attr_from_llfn(llfn: &Value, name: &str) {
llvm::RemoveStringAttrFromFn(llfn, name);
}

/// Get LLVM attribute for the provided inline heuristic.
#[inline]
fn inline_attr<'ll>(cx: &CodegenCx<'ll, '_>, inline: InlineAttr) -> Option<&'ll Attribute> {
Expand Down
28 changes: 1 addition & 27 deletions compiler/rustc_codegen_llvm/src/back/lto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,8 @@ use crate::back::write::{
self, CodegenDiagnosticsStage, DiagnosticHandlers, bitcode_section_name, save_temp_bitcode,
};
use crate::errors::{LlvmError, LtoBitcodeFromRlib};
use crate::llvm::AttributePlace::Function;
use crate::llvm::{self, build_string};
use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx, attributes};
use crate::{LlvmCodegenBackend, ModuleLlvm, SimpleCx};

/// We keep track of the computed LTO cache keys from the previous
/// session to determine which CGUs we can reuse.
Expand Down Expand Up @@ -593,31 +592,6 @@ pub(crate) fn run_pass_manager(
}

if cfg!(llvm_enzyme) && enable_ad && !thin {
let cx =
SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);

for function in cx.get_functions() {
let enzyme_marker = "enzyme_marker";
if attributes::has_string_attr(function, enzyme_marker) {
// Sanity check: Ensure 'noinline' is present before replacing it.
assert!(
attributes::has_attr(function, Function, llvm::AttributeKind::NoInline),
"Expected __enzyme function to have 'noinline' before adding 'alwaysinline'"
);

attributes::remove_from_llfn(function, Function, llvm::AttributeKind::NoInline);
attributes::remove_string_attr_from_llfn(function, enzyme_marker);

assert!(
!attributes::has_string_attr(function, enzyme_marker),
"Expected function to not have 'enzyme_marker'"
);

let always_inline = llvm::AttributeKind::AlwaysInline.create_attr(cx.llcx);
attributes::apply_to_llfn(function, Function, &[always_inline]);
}
}

let opt_stage = llvm::OptStage::FatLTO;
let stage = write::AutodiffStage::PostAD;
if !config.autodiff.contains(&config::AutoDiff::NoPostopt) {
Expand Down
Loading
Loading