diff --git a/compiler/rustc_ast/src/expand/autodiff_attrs.rs b/compiler/rustc_ast/src/expand/autodiff_attrs.rs index 33451f9974835..90f15753e99c9 100644 --- a/compiler/rustc_ast/src/expand/autodiff_attrs.rs +++ b/compiler/rustc_ast/src/expand/autodiff_attrs.rs @@ -6,6 +6,7 @@ use std::fmt::{self, Display, Formatter}; use std::str::FromStr; +use crate::expand::typetree::TypeTree; use crate::expand::{Decodable, Encodable, HashStable_Generic}; use crate::{Ty, TyKind}; @@ -84,6 +85,8 @@ pub struct AutoDiffItem { /// The name of the function being generated pub target: String, pub attrs: AutoDiffAttrs, + pub inputs: Vec, + pub output: TypeTree, } #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)] @@ -275,14 +278,22 @@ impl AutoDiffAttrs { !matches!(self.mode, DiffMode::Error | DiffMode::Source) } - pub fn into_item(self, source: String, target: String) -> AutoDiffItem { - AutoDiffItem { source, target, attrs: self } + pub fn into_item( + self, + source: String, + target: String, + inputs: Vec, + output: TypeTree, + ) -> AutoDiffItem { + AutoDiffItem { source, target, inputs, output, attrs: self } } } impl fmt::Display for AutoDiffItem { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "Differentiating {} -> {}", self.source, self.target)?; - write!(f, " with attributes: {:?}", self.attrs) + write!(f, " with attributes: {:?}", self.attrs)?; + write!(f, " with inputs: {:?}", self.inputs)?; + write!(f, " with output: {:?}", self.output) } } diff --git a/compiler/rustc_ast/src/expand/typetree.rs b/compiler/rustc_ast/src/expand/typetree.rs index 9a2dd2e85e0d6..e7b4f3aff413a 100644 --- a/compiler/rustc_ast/src/expand/typetree.rs +++ b/compiler/rustc_ast/src/expand/typetree.rs @@ -31,6 +31,7 @@ pub enum Kind { Half, Float, Double, + F128, Unknown, } diff --git a/compiler/rustc_codegen_gcc/src/builder.rs b/compiler/rustc_codegen_gcc/src/builder.rs index f7a7a3f8c7e35..5657620879ca1 100644 --- a/compiler/rustc_codegen_gcc/src/builder.rs +++ b/compiler/rustc_codegen_gcc/src/builder.rs @@ -1383,6 +1383,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> { _src_align: Align, size: RValue<'gcc>, flags: MemFlags, + _tt: Option, // Autodiff TypeTrees are LLVM-only, ignored in GCC backend ) { assert!(!flags.contains(MemFlags::NONTEMPORAL), "non-temporal memcpy not supported"); let size = self.intcast(size, self.type_size_t(), false); diff --git a/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs b/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs index eb0a5336a1f13..6aeb9981d5a6e 100644 --- a/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs +++ b/compiler/rustc_codegen_gcc/src/intrinsic/mod.rs @@ -771,6 +771,7 @@ impl<'gcc, 'tcx> ArgAbiExt<'gcc, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> { scratch_align, bx.const_usize(self.layout.size.bytes()), MemFlags::empty(), + None, ); bx.lifetime_end(scratch, scratch_size); diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs index 399f8b6e76287..627d434d120c1 100644 --- a/compiler/rustc_codegen_llvm/src/abi.rs +++ b/compiler/rustc_codegen_llvm/src/abi.rs @@ -244,6 +244,7 @@ impl<'ll, 'tcx> ArgAbiExt<'ll, 'tcx> for ArgAbi<'tcx, Ty<'tcx>> { scratch_align, bx.const_usize(copy_bytes), MemFlags::empty(), + None, ); bx.lifetime_end(llscratch, scratch_size); } diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs index 853d0295238e6..6e88c35a49ac6 100644 --- a/compiler/rustc_codegen_llvm/src/back/lto.rs +++ b/compiler/rustc_codegen_llvm/src/back/lto.rs @@ -537,6 +537,8 @@ fn enable_autodiff_settings(ad: &[config::AutoDiff]) { config::AutoDiff::Enable => {} // We handle this below config::AutoDiff::NoPostopt => {} + // Disables TypeTree generation + config::AutoDiff::NoTT => {} } } // This helps with handling enums for now. diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 427c75d40e9c6..18e1d3b61f9e8 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -2,6 +2,7 @@ use std::borrow::{Borrow, Cow}; use std::ops::Deref; use std::{iter, ptr}; +use rustc_ast::expand::typetree::FncTree; pub(crate) mod autodiff; pub(crate) mod gpu_offload; @@ -1118,11 +1119,12 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { src_align: Align, size: &'ll Value, flags: MemFlags, + tt: Option, ) { assert!(!flags.contains(MemFlags::NONTEMPORAL), "non-temporal memcpy not supported"); let size = self.intcast(size, self.type_isize(), false); let is_volatile = flags.contains(MemFlags::VOLATILE); - unsafe { + let memcpy = unsafe { llvm::LLVMRustBuildMemCpy( self.llbuilder, dst, @@ -1131,7 +1133,16 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { src_align.bytes() as c_uint, size, is_volatile, - ); + ) + }; + + // TypeTree metadata for memcpy is especially important: when Enzyme encounters + // a memcpy during autodiff, it needs to know the structure of the data being + // copied to properly track derivatives. For example, copying an array of floats + // vs. copying a struct with mixed types requires different derivative handling. + // The TypeTree tells Enzyme exactly what memory layout to expect. + if let Some(tt) = tt { + crate::typetree::add_tt(self.cx().llmod, self.cx().llcx, memcpy, tt); } } diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index e2df3265f6f7d..5ec3bbbacb3b9 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -1,6 +1,7 @@ use std::ptr; use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, DiffActivity, DiffMode}; +use rustc_ast::expand::typetree::FncTree; use rustc_codegen_ssa::common::TypeKind; use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods}; use rustc_middle::ty::{PseudoCanonicalInput, Ty, TyCtxt, TypingEnv}; @@ -254,6 +255,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( fn_args: &[&'ll Value], attrs: AutoDiffAttrs, dest: PlaceRef<'tcx, &'ll Value>, + fnc_tree: FncTree, ) { // We have to pick the name depending on whether we want forward or reverse mode autodiff. let mut ad_name: String = match attrs.mode { @@ -330,6 +332,10 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( fn_args, ); + if !fnc_tree.args.is_empty() || !fnc_tree.ret.0.is_empty() { + crate::typetree::add_tt(cx.llmod, cx.llcx, fn_to_diff, fnc_tree); + } + let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None); builder.store_to_place(call, dest.val); diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 49d3dedbeabdf..ded453785af7f 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1202,6 +1202,9 @@ fn codegen_autodiff<'ll, 'tcx>( &mut diff_attrs.input_activity, ); + let fnc_tree = + rustc_middle::ty::fnc_typetrees(tcx, fn_source.ty(tcx, TypingEnv::fully_monomorphized())); + // Build body generate_enzyme_call( bx, @@ -1212,6 +1215,7 @@ fn codegen_autodiff<'ll, 'tcx>( &val_arr, diff_attrs.clone(), result, + fnc_tree, ); } diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs index 0fcf31d799308..feafe130b445d 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -67,6 +67,7 @@ mod llvm_util; mod mono_item; mod type_; mod type_of; +mod typetree; mod va_arg; mod value; diff --git a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs index 56d756e52cce1..348004535cc08 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs @@ -3,9 +3,36 @@ use libc::{c_char, c_uint}; use super::MetadataKindId; -use super::ffi::{AttributeKind, BasicBlock, Metadata, Module, Type, Value}; +use super::ffi::{AttributeKind, BasicBlock, Context, Metadata, Module, Type, Value}; use crate::llvm::{Bool, Builder}; +// TypeTree types +pub(crate) type CTypeTreeRef = *mut EnzymeTypeTree; + +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub(crate) struct EnzymeTypeTree { + _unused: [u8; 0], +} + +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +#[allow(non_camel_case_types)] +pub(crate) enum CConcreteType { + DT_Anything = 0, + DT_Integer = 1, + DT_Pointer = 2, + DT_Half = 3, + DT_Float = 4, + DT_Double = 5, + DT_Unknown = 6, + DT_FP128 = 9, +} + +pub(crate) struct TypeTree { + pub(crate) inner: CTypeTreeRef, +} + #[link(name = "llvm-wrapper", kind = "static")] unsafe extern "C" { // Enzyme @@ -68,10 +95,33 @@ pub(crate) mod Enzyme_AD { use libc::c_void; + use super::{CConcreteType, CTypeTreeRef, Context}; + unsafe extern "C" { pub(crate) fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8); pub(crate) fn EnzymeSetCLString(arg1: *mut ::std::os::raw::c_void, arg2: *const c_char); } + + // TypeTree functions + unsafe extern "C" { + pub(crate) fn EnzymeNewTypeTree() -> CTypeTreeRef; + pub(crate) fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef; + pub(crate) fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef; + pub(crate) fn EnzymeFreeTypeTree(CTT: CTypeTreeRef); + pub(crate) fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef) -> bool; + pub(crate) fn EnzymeTypeTreeOnlyEq(arg1: CTypeTreeRef, pos: i64); + pub(crate) fn EnzymeTypeTreeData0Eq(arg1: CTypeTreeRef); + pub(crate) fn EnzymeTypeTreeShiftIndiciesEq( + arg1: CTypeTreeRef, + data_layout: *const c_char, + offset: i64, + max_size: i64, + add_offset: u64, + ); + pub(crate) fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char; + pub(crate) fn EnzymeTypeTreeToStringFree(arg1: *const c_char); + } + unsafe extern "C" { static mut EnzymePrintPerf: c_void; static mut EnzymePrintActivity: c_void; @@ -141,6 +191,57 @@ pub(crate) use self::Fallback_AD::*; pub(crate) mod Fallback_AD { #![allow(unused_variables)] + use libc::c_char; + + use super::{CConcreteType, CTypeTreeRef, Context}; + + // TypeTree function fallbacks + pub(crate) unsafe fn EnzymeNewTypeTree() -> CTypeTreeRef { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeFreeTypeTree(CTT: CTypeTreeRef) { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef) -> bool { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeTypeTreeOnlyEq(arg1: CTypeTreeRef, pos: i64) { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeTypeTreeData0Eq(arg1: CTypeTreeRef) { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeTypeTreeShiftIndiciesEq( + arg1: CTypeTreeRef, + data_layout: *const c_char, + offset: i64, + max_size: i64, + add_offset: u64, + ) { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char { + unimplemented!() + } + + pub(crate) unsafe fn EnzymeTypeTreeToStringFree(arg1: *const c_char) { + unimplemented!() + } + pub(crate) fn set_inline(val: bool) { unimplemented!() } @@ -169,3 +270,83 @@ pub(crate) mod Fallback_AD { unimplemented!() } } + +impl TypeTree { + pub(crate) fn new() -> TypeTree { + let inner = unsafe { EnzymeNewTypeTree() }; + TypeTree { inner } + } + + pub(crate) fn from_type(t: CConcreteType, ctx: &Context) -> TypeTree { + let inner = unsafe { EnzymeNewTypeTreeCT(t, ctx) }; + TypeTree { inner } + } + + pub(crate) fn merge(self, other: Self) -> Self { + unsafe { + EnzymeMergeTypeTree(self.inner, other.inner); + } + drop(other); + self + } + + #[must_use] + pub(crate) fn shift( + self, + layout: &str, + offset: isize, + max_size: isize, + add_offset: usize, + ) -> Self { + let layout = std::ffi::CString::new(layout).unwrap(); + + unsafe { + EnzymeTypeTreeShiftIndiciesEq( + self.inner, + layout.as_ptr(), + offset as i64, + max_size as i64, + add_offset as u64, + ); + } + + self + } +} + +impl Clone for TypeTree { + fn clone(&self) -> Self { + let inner = unsafe { EnzymeNewTypeTreeTR(self.inner) }; + TypeTree { inner } + } +} + +impl std::fmt::Display for TypeTree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let ptr = unsafe { EnzymeTypeTreeToString(self.inner) }; + let cstr = unsafe { std::ffi::CStr::from_ptr(ptr) }; + match cstr.to_str() { + Ok(x) => write!(f, "{}", x)?, + Err(err) => write!(f, "could not parse: {}", err)?, + } + + // delete C string pointer + unsafe { + EnzymeTypeTreeToStringFree(ptr); + } + + Ok(()) + } +} + +impl std::fmt::Debug for TypeTree { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + ::fmt(self, f) + } +} + +impl Drop for TypeTree { + fn drop(&mut self) { + unsafe { EnzymeFreeTypeTree(self.inner) } + } +} diff --git a/compiler/rustc_codegen_llvm/src/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs new file mode 100644 index 0000000000000..8c0d255bba84e --- /dev/null +++ b/compiler/rustc_codegen_llvm/src/typetree.rs @@ -0,0 +1,138 @@ +use rustc_ast::expand::typetree::FncTree; +#[cfg(llvm_enzyme)] +use { + crate::attributes, + rustc_ast::expand::typetree::TypeTree as RustTypeTree, + std::ffi::{CString, c_char, c_uint}, +}; + +use crate::llvm::{self, Value}; + +/// Converts a Rust TypeTree to Enzyme's internal TypeTree format +/// +/// This function takes a Rust-side TypeTree (from rustc_ast::expand::typetree) +/// and converts it to Enzyme's internal C++ TypeTree representation that +/// Enzyme can understand during differentiation analysis. +#[cfg(llvm_enzyme)] +fn to_enzyme_typetree( + rust_typetree: RustTypeTree, + data_layout: &str, + llcx: &llvm::Context, +) -> llvm::TypeTree { + // Start with an empty TypeTree + let mut enzyme_tt = llvm::TypeTree::new(); + + // Convert each Type in the Rust TypeTree to Enzyme format + for rust_type in rust_typetree.0 { + let concrete_type = match rust_type.kind { + rustc_ast::expand::typetree::Kind::Anything => llvm::CConcreteType::DT_Anything, + rustc_ast::expand::typetree::Kind::Integer => llvm::CConcreteType::DT_Integer, + rustc_ast::expand::typetree::Kind::Pointer => llvm::CConcreteType::DT_Pointer, + rustc_ast::expand::typetree::Kind::Half => llvm::CConcreteType::DT_Half, + rustc_ast::expand::typetree::Kind::Float => llvm::CConcreteType::DT_Float, + rustc_ast::expand::typetree::Kind::Double => llvm::CConcreteType::DT_Double, + rustc_ast::expand::typetree::Kind::F128 => llvm::CConcreteType::DT_FP128, + rustc_ast::expand::typetree::Kind::Unknown => llvm::CConcreteType::DT_Unknown, + }; + + // Create a TypeTree for this specific type + let type_tt = llvm::TypeTree::from_type(concrete_type, llcx); + + // Apply offset if specified + let type_tt = if rust_type.offset == -1 { + type_tt // -1 means everywhere/no specific offset + } else { + // Apply specific offset positioning + type_tt.shift(data_layout, rust_type.offset, rust_type.size as isize, 0) + }; + + // Merge this type into the main TypeTree + enzyme_tt = enzyme_tt.merge(type_tt); + } + + enzyme_tt +} + +// Attaches TypeTree information to LLVM function as enzyme_type attributes. +#[cfg(llvm_enzyme)] +pub(crate) fn add_tt<'ll>( + llmod: &'ll llvm::Module, + llcx: &'ll llvm::Context, + fn_def: &'ll Value, + tt: FncTree, +) { + let inputs = tt.args; + let ret_tt: RustTypeTree = tt.ret; + + // Get LLVM data layout string for TypeTree conversion + let llvm_data_layout: *const c_char = unsafe { llvm::LLVMGetDataLayoutStr(&*llmod) }; + let llvm_data_layout = + std::str::from_utf8(unsafe { std::ffi::CStr::from_ptr(llvm_data_layout) }.to_bytes()) + .expect("got a non-UTF8 data-layout from LLVM"); + + // Attribute name that Enzyme recognizes for TypeTree information + let attr_name = "enzyme_type"; + let c_attr_name = CString::new(attr_name).unwrap(); + + // Attach TypeTree attributes to each input parameter + // Enzyme uses these to understand parameter memory layouts during differentiation + for (i, input) in inputs.iter().enumerate() { + unsafe { + // Convert Rust TypeTree to Enzyme's internal format + let enzyme_tt = to_enzyme_typetree(input.clone(), llvm_data_layout, llcx); + + // Serialize TypeTree to string format that Enzyme can parse + let c_str = llvm::EnzymeTypeTreeToString(enzyme_tt.inner); + let c_str = std::ffi::CStr::from_ptr(c_str); + + // Create LLVM string attribute with TypeTree information + let attr = llvm::LLVMCreateStringAttribute( + llcx, + c_attr_name.as_ptr(), + c_attr_name.as_bytes().len() as c_uint, + c_str.as_ptr(), + c_str.to_bytes().len() as c_uint, + ); + + // Attach attribute to the specific function parameter + // Note: ArgumentPlace uses 0-based indexing, but LLVM uses 1-based for arguments + attributes::apply_to_llfn(fn_def, llvm::AttributePlace::Argument(i as u32), &[attr]); + + // Free the C string to prevent memory leaks + llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr()); + } + } + + // Attach TypeTree attribute to the return type + // Enzyme needs this to understand how to handle return value derivatives + unsafe { + let enzyme_tt = to_enzyme_typetree(ret_tt, llvm_data_layout, llcx); + let c_str = llvm::EnzymeTypeTreeToString(enzyme_tt.inner); + let c_str = std::ffi::CStr::from_ptr(c_str); + + let ret_attr = llvm::LLVMCreateStringAttribute( + llcx, + c_attr_name.as_ptr(), + c_attr_name.as_bytes().len() as c_uint, + c_str.as_ptr(), + c_str.to_bytes().len() as c_uint, + ); + + // Attach to function return type + attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]); + + // Free the C string + llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr()); + } +} + +// Fallback implementation when Enzyme is not available +#[cfg(not(llvm_enzyme))] +pub(crate) fn add_tt<'ll>( + _llmod: &'ll llvm::Module, + _llcx: &'ll llvm::Context, + _fn_def: &'ll Value, + _tt: FncTree, +) { + unimplemented!() +} diff --git a/compiler/rustc_codegen_llvm/src/va_arg.rs b/compiler/rustc_codegen_llvm/src/va_arg.rs index ce079f3cb0af1..99d72b63a10ce 100644 --- a/compiler/rustc_codegen_llvm/src/va_arg.rs +++ b/compiler/rustc_codegen_llvm/src/va_arg.rs @@ -735,6 +735,7 @@ fn copy_to_temporary_if_more_aligned<'ll, 'tcx>( src_align, bx.const_u32(layout.layout.size().bytes() as u32), MemFlags::empty(), + None, ); tmp } else { diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index c3dc3e42b83d8..514b490e52fa1 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -1623,6 +1623,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { align, bx.const_usize(copy_bytes), MemFlags::empty(), + None, ); // ...and then load it with the ABI type. llval = load_cast(bx, cast, llscratch, scratch_align); diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs index 3c667b8e88203..befa00c6861ed 100644 --- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs +++ b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs @@ -30,7 +30,7 @@ fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( if allow_overlap { bx.memmove(dst, align, src, align, size, flags); } else { - bx.memcpy(dst, align, src, align, size, flags); + bx.memcpy(dst, align, src, align, size, flags, None); } } diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs index f164e0f912373..0a50d7f18dbef 100644 --- a/compiler/rustc_codegen_ssa/src/mir/statement.rs +++ b/compiler/rustc_codegen_ssa/src/mir/statement.rs @@ -90,7 +90,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { let align = pointee_layout.align; let dst = dst_val.immediate(); let src = src_val.immediate(); - bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty()); + bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), None); } mir::StatementKind::FakeRead(..) | mir::StatementKind::Retag { .. } diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs index f417d1a7bf724..da1456973779b 100644 --- a/compiler/rustc_codegen_ssa/src/traits/builder.rs +++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs @@ -424,6 +424,7 @@ pub trait BuilderMethods<'a, 'tcx>: src_align: Align, size: Self::Value, flags: MemFlags, + tt: Option, ); fn memmove( &mut self, @@ -480,7 +481,7 @@ pub trait BuilderMethods<'a, 'tcx>: temp.val.store_with_flags(self, dst.with_type(layout), flags); } else if !layout.is_zst() { let bytes = self.const_usize(layout.size.bytes()); - self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags); + self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags, None); } } diff --git a/compiler/rustc_interface/src/tests.rs b/compiler/rustc_interface/src/tests.rs index 4425877308a78..32fccd4bf7c2c 100644 --- a/compiler/rustc_interface/src/tests.rs +++ b/compiler/rustc_interface/src/tests.rs @@ -764,7 +764,7 @@ fn test_unstable_options_tracking_hash() { tracked!(allow_features, Some(vec![String::from("lang_items")])); tracked!(always_encode_mir, true); tracked!(assume_incomplete_release, true); - tracked!(autodiff, vec![AutoDiff::Enable]); + tracked!(autodiff, vec![AutoDiff::Enable, AutoDiff::NoTT]); tracked!(binary_dep_depinfo, true); tracked!(box_noalias, false); tracked!( diff --git a/compiler/rustc_middle/src/error.rs b/compiler/rustc_middle/src/error.rs index 7520bc262c6bc..e20c65efdec75 100644 --- a/compiler/rustc_middle/src/error.rs +++ b/compiler/rustc_middle/src/error.rs @@ -37,7 +37,6 @@ pub(crate) struct OpaqueHiddenTypeMismatch<'tcx> { pub sub: TypeMismatchReason, } -// FIXME(autodiff): I should get used somewhere #[derive(Diagnostic)] #[diag(middle_unsupported_union)] pub struct UnsupportedUnion { diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs index a7298af502ef1..391d1ea9896d0 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -25,6 +25,7 @@ pub use generic_args::{GenericArgKind, TermKind, *}; pub use generics::*; pub use intrinsic::IntrinsicDef; use rustc_abi::{Align, FieldIdx, Integer, IntegerType, ReprFlags, ReprOptions, VariantIdx}; +use rustc_ast::expand::typetree::{FncTree, Kind, Type, TypeTree}; use rustc_ast::node_id::NodeMap; pub use rustc_ast_ir::{Movability, Mutability, try_visit}; use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap, FxIndexSet}; @@ -2222,3 +2223,122 @@ pub struct DestructuredConst<'tcx> { pub variant: Option, pub fields: &'tcx [ty::Const<'tcx>], } + +/// Generate TypeTree information for autodiff. +/// This function creates TypeTree metadata that describes the memory layout +/// of function parameters and return types for Enzyme autodiff. +pub fn fnc_typetrees<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>) -> FncTree { + // Check if TypeTrees are disabled via NoTT flag + if tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::NoTT) { + return FncTree { args: vec![], ret: TypeTree::new() }; + } + + // Check if this is actually a function type + if !fn_ty.is_fn() { + return FncTree { args: vec![], ret: TypeTree::new() }; + } + + // Get the function signature + let fn_sig = fn_ty.fn_sig(tcx); + let sig = tcx.instantiate_bound_regions_with_erased(fn_sig); + + // Create TypeTrees for each input parameter + let mut args = vec![]; + for ty in sig.inputs().iter() { + let type_tree = typetree_from_ty(tcx, *ty); + args.push(type_tree); + } + + // Create TypeTree for return type + let ret = typetree_from_ty(tcx, sig.output()); + + FncTree { args, ret } +} + +/// Generate TypeTree for a specific type. +/// This function analyzes a Rust type and creates appropriate TypeTree metadata. +pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree { + if ty.is_scalar() { + let (kind, size) = if ty.is_integral() || ty.is_char() || ty.is_bool() { + (Kind::Integer, ty.primitive_size(tcx).bytes_usize()) + } else if ty.is_floating_point() { + match ty { + x if x == tcx.types.f16 => (Kind::Half, 2), + x if x == tcx.types.f32 => (Kind::Float, 4), + x if x == tcx.types.f64 => (Kind::Double, 8), + x if x == tcx.types.f128 => (Kind::F128, 16), + _ => return TypeTree::new(), + } + } else { + return TypeTree::new(); + }; + + return TypeTree(vec![Type { offset: -1, size, kind, child: TypeTree::new() }]); + } + + if ty.is_ref() || ty.is_raw_ptr() || ty.is_box() { + let inner_ty = if let Some(inner) = ty.builtin_deref(true) { + inner + } else { + return TypeTree::new(); + }; + + let child = typetree_from_ty(tcx, inner_ty); + return TypeTree(vec![Type { + offset: -1, + size: tcx.data_layout.pointer_size().bytes_usize(), + kind: Kind::Pointer, + child, + }]); + } + + if ty.is_array() { + if let ty::Array(element_ty, len_const) = ty.kind() { + let len = len_const.try_to_target_usize(tcx).unwrap_or(0); + if len == 0 { + return TypeTree::new(); + } + + let element_tree = typetree_from_ty(tcx, *element_ty); + + let element_layout = tcx + .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(*element_ty)) + .ok() + .map(|layout| layout.size.bytes_usize()) + .unwrap_or(0); + + if element_layout == 0 { + return TypeTree::new(); + } + + let mut types = Vec::new(); + for i in 0..len { + let base_offset = (i as usize * element_layout) as isize; + + for elem_type in &element_tree.0 { + types.push(Type { + offset: if elem_type.offset == -1 { + base_offset + } else { + base_offset + elem_type.offset + }, + size: elem_type.size, + kind: elem_type.kind, + child: elem_type.child.clone(), + }); + } + } + + return TypeTree(types); + } + } + + if ty.is_slice() { + if let ty::Slice(element_ty) = ty.kind() { + let element_tree = typetree_from_ty(tcx, *element_ty); + return element_tree; + } + } + + TypeTree::new() +} diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs index 9793d8091e2ab..f13304b67c45d 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -256,6 +256,8 @@ pub enum AutoDiff { LooseTypes, /// Runs Enzyme's aggressive inlining Inline, + /// Disable Type Tree + NoTT, } /// Settings for `-Z instrument-xray` flag. diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs index 6d5be2d92cd2a..e83a022dd2246 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -725,7 +725,7 @@ mod desc { pub(crate) const parse_list: &str = "a space-separated list of strings"; pub(crate) const parse_list_with_polarity: &str = "a comma-separated list of strings, with elements beginning with + or -"; - pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintTAFn`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `PrintModFinal`, `PrintPasses`, `NoPostopt`, `LooseTypes`, `Inline`"; + pub(crate) const parse_autodiff: &str = "a comma separated list of settings: `Enable`, `PrintSteps`, `PrintTA`, `PrintTAFn`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `PrintModFinal`, `PrintPasses`, `NoPostopt`, `LooseTypes`, `Inline`, `NoTT`"; pub(crate) const parse_offload: &str = "a comma separated list of settings: `Enable`"; pub(crate) const parse_comma_list: &str = "a comma-separated list of strings"; pub(crate) const parse_opt_comma_list: &str = parse_comma_list; @@ -1412,6 +1412,7 @@ pub mod parse { "PrintPasses" => AutoDiff::PrintPasses, "LooseTypes" => AutoDiff::LooseTypes, "Inline" => AutoDiff::Inline, + "NoTT" => AutoDiff::NoTT, _ => { // FIXME(ZuseZ4): print an error saying which value is not recognized return false; diff --git a/src/tools/enzyme b/src/tools/enzyme index 58af4e9e6c047..216fbf2078ff0 160000 --- a/src/tools/enzyme +++ b/src/tools/enzyme @@ -1 +1 @@ -Subproject commit 58af4e9e6c047534ba059b12af17cecb8a2e9f9e +Subproject commit 216fbf2078ff0581d7b33bdd90a5d93762fc81e7 diff --git a/tests/codegen-llvm/autodiff/typetree.rs b/tests/codegen-llvm/autodiff/typetree.rs new file mode 100644 index 0000000000000..1cb0c2fb68be3 --- /dev/null +++ b/tests/codegen-llvm/autodiff/typetree.rs @@ -0,0 +1,33 @@ +//@ compile-flags: -Zautodiff=Enable -C opt-level=3 -Clto=fat +//@ no-prefer-dynamic +//@ needs-enzyme + +// Test that basic autodiff still works with our TypeTree infrastructure +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_simple, Duplicated, Active)] +#[no_mangle] +#[inline(never)] +fn simple(x: &f64) -> f64 { + 2.0 * x +} + +// CHECK-LABEL: @simple +// CHECK: fmul double + +// The derivative function should be generated normally +// CHECK-LABEL: diffesimple +// CHECK: fadd fast double + +fn main() { + let x = std::hint::black_box(3.0); + let output = simple(&x); + assert_eq!(6.0, output); + + let mut df_dx = 0.0; + let output_ = d_simple(&x, &mut df_dx, 1.0); + assert_eq!(output, output_); + assert_eq!(2.0, df_dx); +} diff --git a/tests/run-make/autodiff/type-trees/array-typetree/array.check b/tests/run-make/autodiff/type-trees/array-typetree/array.check new file mode 100644 index 0000000000000..7513458b8ab81 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/array-typetree/array.check @@ -0,0 +1,4 @@ +; Check that array TypeTree metadata is correctly generated +; Should show Float@double at each array element offset (0, 8, 16, 24, 32 bytes) + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_array{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/array-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/array-typetree/rmake.rs new file mode 100644 index 0000000000000..20b6a06690625 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/array-typetree/rmake.rs @@ -0,0 +1,9 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + llvm_filecheck().patterns("array.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/array-typetree/test.rs b/tests/run-make/autodiff/type-trees/array-typetree/test.rs new file mode 100644 index 0000000000000..f54ebf5a4c7b5 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/array-typetree/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_array(arr: &[f64; 5]) -> f64 { + arr[0] + arr[1] + arr[2] + arr[3] + arr[4] +} + +fn main() { + let arr = [1.0, 2.0, 3.0, 4.0, 5.0]; + let mut d_arr = [0.0; 5]; + let _result = d_test(&arr, &mut d_arr, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy-ir.check b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy-ir.check new file mode 100644 index 0000000000000..3a59f06dda1e4 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy-ir.check @@ -0,0 +1,8 @@ +; Check that enzyme_type attributes are present in the LLVM IR function definition +; This verifies our TypeTree system correctly attaches metadata for Enzyme + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_memcpy({{.*}}"enzyme_type"="{[]:Pointer}" + +; Check that llvm.memcpy exists (either call or declare) +CHECK: {{(call|declare).*}}@llvm.memcpy + diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.check b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.check new file mode 100644 index 0000000000000..ae70830297a78 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.check @@ -0,0 +1,13 @@ +CHECK: force_memcpy + +CHECK: @llvm.memcpy.p0.p0.i64 + +CHECK: test_memcpy - {[-1]:Float@double} |{[-1]:Pointer}:{} + +CHECK-DAG: ptr %{{[0-9]+}}: {[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Float@double, [-1,16]:Float@double, [-1,24]:Float@double} + +CHECK-DAG: load double{{.*}}: {[-1]:Float@double} + +CHECK-DAG: fmul double{{.*}}: {[-1]:Float@double} + +CHECK-DAG: fadd double{{.*}}: {[-1]:Float@double} \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs new file mode 100644 index 0000000000000..3c1029190c88a --- /dev/null +++ b/tests/run-make/autodiff/type-trees/memcpy-typetree/memcpy.rs @@ -0,0 +1,36 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; +use std::ptr; + +#[inline(never)] +fn force_memcpy(src: *const f64, dst: *mut f64, count: usize) { + unsafe { + ptr::copy_nonoverlapping(src, dst, count); + } +} + +#[autodiff_reverse(d_test_memcpy, Duplicated, Active)] +#[no_mangle] +fn test_memcpy(input: &[f64; 128]) -> f64 { + let mut local_data = [0.0f64; 128]; + + // Use a separate function to prevent inlining and optimization + force_memcpy(input.as_ptr(), local_data.as_mut_ptr(), 128); + + // Sum only first few elements to keep the computation simple + local_data[0] * local_data[0] + + local_data[1] * local_data[1] + + local_data[2] * local_data[2] + + local_data[3] * local_data[3] +} + +fn main() { + let input = [1.0; 128]; + let mut d_input = [0.0; 128]; + let result = test_memcpy(&input); + let result_d = d_test_memcpy(&input, &mut d_input, 1.0); + + assert_eq!(result, result_d); + println!("Memcpy test passed: result = {}", result); +} diff --git a/tests/run-make/autodiff/type-trees/memcpy-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/memcpy-typetree/rmake.rs new file mode 100644 index 0000000000000..b4c650330fe94 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/memcpy-typetree/rmake.rs @@ -0,0 +1,39 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // First, compile to LLVM IR to check for enzyme_type attributes + let _ir_output = rustc() + .input("memcpy.rs") + .arg("-Zautodiff=Enable") + .arg("-Zautodiff=NoPostopt") + .opt_level("0") + .arg("--emit=llvm-ir") + .arg("-o") + .arg("main.ll") + .run(); + + // Then compile with TypeTree analysis output for the existing checks + let output = rustc() + .input("memcpy.rs") + .arg("-Zautodiff=Enable,PrintTAFn=test_memcpy") + .arg("-Zautodiff=NoPostopt") + .opt_level("3") + .arg("-Clto=fat") + .arg("-g") + .run(); + + let stdout = output.stdout_utf8(); + let stderr = output.stderr_utf8(); + let ir_content = rfs::read_to_string("main.ll"); + + rfs::write("memcpy.stdout", &stdout); + rfs::write("memcpy.stderr", &stderr); + rfs::write("main.ir", &ir_content); + + llvm_filecheck().patterns("memcpy.check").stdin_buf(stdout).run(); + + llvm_filecheck().patterns("memcpy-ir.check").stdin_buf(ir_content).run(); +} diff --git a/tests/run-make/autodiff/type-trees/nott-flag/nott.check b/tests/run-make/autodiff/type-trees/nott-flag/nott.check new file mode 100644 index 0000000000000..8d23e2ee31957 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/nott-flag/nott.check @@ -0,0 +1,5 @@ +// Check that enzyme_type attributes are NOT present when NoTT flag is used +// This verifies the NoTT flag correctly disables TypeTree metadata + +CHECK: define{{.*}}@square +CHECK-NOT: "enzyme_type" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/nott-flag/rmake.rs b/tests/run-make/autodiff/type-trees/nott-flag/rmake.rs new file mode 100644 index 0000000000000..de540b990cabd --- /dev/null +++ b/tests/run-make/autodiff/type-trees/nott-flag/rmake.rs @@ -0,0 +1,30 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Test with NoTT flag - should not generate TypeTree metadata + rustc() + .input("test.rs") + .arg("-Zautodiff=Enable,NoTT") + .emit("llvm-ir") + .arg("-o") + .arg("nott.ll") + .run(); + + // Test without NoTT flag - should generate TypeTree metadata + rustc() + .input("test.rs") + .arg("-Zautodiff=Enable") + .emit("llvm-ir") + .arg("-o") + .arg("with_tt.ll") + .run(); + + // Verify NoTT version does NOT have enzyme_type attributes + llvm_filecheck().patterns("nott.check").stdin_buf(rfs::read("nott.ll")).run(); + + // Verify TypeTree version DOES have enzyme_type attributes + llvm_filecheck().patterns("with_tt.check").stdin_buf(rfs::read("with_tt.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/nott-flag/test.rs b/tests/run-make/autodiff/type-trees/nott-flag/test.rs new file mode 100644 index 0000000000000..de3549c37c679 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/nott-flag/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_square, Duplicated, Active)] +#[no_mangle] +fn square(x: &f64) -> f64 { + x * x +} + +fn main() { + let x = 2.0; + let mut dx = 0.0; + let _result = d_square(&x, &mut dx, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/nott-flag/with_tt.check b/tests/run-make/autodiff/type-trees/nott-flag/with_tt.check new file mode 100644 index 0000000000000..53eafc10a65da --- /dev/null +++ b/tests/run-make/autodiff/type-trees/nott-flag/with_tt.check @@ -0,0 +1,4 @@ +// Check that enzyme_type attributes are present when TypeTree is enabled +// This verifies our TypeTree metadata attachment is working + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@square{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/f128.check b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/f128.check new file mode 100644 index 0000000000000..31cef113420db --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/f128.check @@ -0,0 +1,4 @@ +; Check that f128 TypeTree metadata is correctly generated +; f128 maps to Unknown in our current implementation since CConcreteType doesn't have DT_F128 + +CHECK: define{{.*}}"enzyme_type"={{.*}}@test_f128{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/rmake.rs new file mode 100644 index 0000000000000..44320ecdd5714 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/rmake.rs @@ -0,0 +1,12 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Compile with TypeTree enabled and emit LLVM IR + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + + // Check that f128 TypeTree metadata is correctly generated + llvm_filecheck().patterns("f128.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/test.rs b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/test.rs new file mode 100644 index 0000000000000..5c71baa3e6999 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f128-typetree/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff, f128)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_f128(x: &f128) -> f128 { + *x * *x +} + +fn main() { + let x = 2.0_f128; + let mut dx = 0.0_f128; + let _result = d_test(&x, &mut dx, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/f16.check b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/f16.check new file mode 100644 index 0000000000000..97a3964a569b1 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/f16.check @@ -0,0 +1,4 @@ +; Check that f16 TypeTree metadata is correctly generated +; Should show Half for f16 values and Pointer for references + +CHECK: define{{.*}}"enzyme_type"={{.*}}@test_f16{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/rmake.rs new file mode 100644 index 0000000000000..0aebdbf55209b --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/rmake.rs @@ -0,0 +1,12 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Compile with TypeTree enabled and emit LLVM IR + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + + // Check that f16 TypeTree metadata is correctly generated + llvm_filecheck().patterns("f16.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/test.rs b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/test.rs new file mode 100644 index 0000000000000..6b68e8252f4c0 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f16-typetree/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff, f16)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_f16(x: &f16) -> f16 { + *x * *x +} + +fn main() { + let x = 2.0_f16; + let mut dx = 0.0_f16; + let _result = d_test(&x, &mut dx, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/f32.check b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/f32.check new file mode 100644 index 0000000000000..b32d5086d07d1 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/f32.check @@ -0,0 +1,4 @@ +; Check that f32 TypeTree metadata is correctly generated +; Should show Float@float for f32 values and Pointer for references + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@float}"{{.*}}@test_f32{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/rmake.rs new file mode 100644 index 0000000000000..ee3ab753bf505 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/rmake.rs @@ -0,0 +1,12 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Compile with TypeTree enabled and emit LLVM IR + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + + // Check that f32 TypeTree metadata is correctly generated + llvm_filecheck().patterns("f32.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/test.rs b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/test.rs new file mode 100644 index 0000000000000..56c118399ee4b --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f32-typetree/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_f32(x: &f32) -> f32 { + x * x +} + +fn main() { + let x = 2.0_f32; + let mut dx = 0.0_f32; + let _result = d_test(&x, &mut dx, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/f64.check b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/f64.check new file mode 100644 index 0000000000000..0e9d9c03a0f86 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/f64.check @@ -0,0 +1,4 @@ +; Check that f64 TypeTree metadata is correctly generated +; Should show Float@double for f64 values and Pointer for references + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_f64{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/rmake.rs new file mode 100644 index 0000000000000..5fac9b23bc80f --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/rmake.rs @@ -0,0 +1,12 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Compile with TypeTree enabled and emit LLVM IR + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + + // Check that f64 TypeTree metadata is correctly generated + llvm_filecheck().patterns("f64.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/test.rs b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/test.rs new file mode 100644 index 0000000000000..235360b76b23d --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/f64-typetree/test.rs @@ -0,0 +1,15 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_f64(x: &f64) -> f64 { + x * x +} + +fn main() { + let x = 2.0_f64; + let mut dx = 0.0_f64; + let _result = d_test(&x, &mut dx, 1.0); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/i32.check b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/i32.check new file mode 100644 index 0000000000000..aa8e5223e8ed5 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/i32.check @@ -0,0 +1,4 @@ +; Check that i32 TypeTree metadata is correctly generated +; Should show Integer for i32 values (integers are typically Const in autodiff) + +CHECK: define{{.*}}"enzyme_type"="{[]:Integer}"{{.*}}@test_i32{{.*}}"enzyme_type"="{[]:Integer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/rmake.rs new file mode 100644 index 0000000000000..a40fd55d88adf --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/rmake.rs @@ -0,0 +1,12 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + // Compile with TypeTree enabled and emit LLVM IR + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + + // Check that i32 TypeTree metadata is correctly generated + llvm_filecheck().patterns("i32.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/test.rs b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/test.rs new file mode 100644 index 0000000000000..530c137f62b9c --- /dev/null +++ b/tests/run-make/autodiff/type-trees/scalar-types/i32-typetree/test.rs @@ -0,0 +1,14 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Const, Active)] +#[no_mangle] +fn test_i32(x: i32) -> i32 { + x * x +} + +fn main() { + let x = 5_i32; + let _result = d_test(x, 1); +} diff --git a/tests/run-make/autodiff/type-trees/slice-typetree/rmake.rs b/tests/run-make/autodiff/type-trees/slice-typetree/rmake.rs new file mode 100644 index 0000000000000..b81fb50bf1a75 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-typetree/rmake.rs @@ -0,0 +1,9 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +fn main() { + rustc().input("test.rs").arg("-Zautodiff=Enable").emit("llvm-ir").run(); + llvm_filecheck().patterns("slice.check").stdin_buf(rfs::read("test.ll")).run(); +} diff --git a/tests/run-make/autodiff/type-trees/slice-typetree/slice.check b/tests/run-make/autodiff/type-trees/slice-typetree/slice.check new file mode 100644 index 0000000000000..8fc9e9c4f1b17 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-typetree/slice.check @@ -0,0 +1,4 @@ +; Check that slice TypeTree metadata is correctly generated +; Should show Float@double for slice elements + +CHECK: define{{.*}}"enzyme_type"="{[]:Float@double}"{{.*}}@test_slice{{.*}}"enzyme_type"="{[]:Pointer}" \ No newline at end of file diff --git a/tests/run-make/autodiff/type-trees/slice-typetree/test.rs b/tests/run-make/autodiff/type-trees/slice-typetree/test.rs new file mode 100644 index 0000000000000..7117fa3844f59 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/slice-typetree/test.rs @@ -0,0 +1,16 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +#[autodiff_reverse(d_test, Duplicated, Active)] +#[no_mangle] +fn test_slice(slice: &[f64]) -> f64 { + slice.iter().sum() +} + +fn main() { + let arr = [1.0, 2.0, 3.0, 4.0, 5.0]; + let slice = &arr[..]; + let mut d_slice = [0.0; 5]; + let _result = d_test(slice, &mut d_slice[..], 1.0); +} diff --git a/tests/ui/autodiff/flag_nott.rs b/tests/ui/autodiff/flag_nott.rs new file mode 100644 index 0000000000000..faa9949fe8167 --- /dev/null +++ b/tests/ui/autodiff/flag_nott.rs @@ -0,0 +1,19 @@ +//@ compile-flags: -Zautodiff=Enable,NoTT +//@ needs-enzyme +//@ check-pass + +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +// Test that NoTT flag is accepted and doesn't cause compilation errors +#[autodiff_reverse(d_square, Duplicated, Active)] +fn square(x: &f64) -> f64 { + x * x +} + +fn main() { + let x = 2.0; + let mut dx = 0.0; + let result = d_square(&x, &mut dx, 1.0); +} diff --git a/typos.toml b/typos.toml index b0ff48f8fa28b..c7d5b0000b9d7 100644 --- a/typos.toml +++ b/typos.toml @@ -53,6 +53,7 @@ ERROR_DS_NOT_AUTHORITIVE_FOR_DST_NC = "ERROR_DS_NOT_AUTHORITIVE_FOR_DST_NC" ERROR_MCA_OCCURED = "ERROR_MCA_OCCURED" ERRNO_ACCES = "ERRNO_ACCES" tolen = "tolen" +EnzymeTypeTreeShiftIndiciesEq = "EnzymeTypeTreeShiftIndiciesEq" [default] extend-ignore-words-re = [