diff --git a/compiler/rustc_ast/src/expand/typetree.rs b/compiler/rustc_ast/src/expand/typetree.rs index 9619c80904426..00fea30ba94e5 100644 --- a/compiler/rustc_ast/src/expand/typetree.rs +++ b/compiler/rustc_ast/src/expand/typetree.rs @@ -28,6 +28,9 @@ pub enum Kind { Anything, Integer, Pointer, + // We prefer to directly lower to things that our Enzyme backend supports. + // However, it's sometimes convenient to pass ptr+int as one type. + RustSlice, Half, Float, Double, @@ -57,6 +60,9 @@ impl TypeTree { } Self(ints) } + pub fn add_indirection(self) -> Self { + Self(vec![Type { offset: 0, size: 1, kind: Kind::Pointer, child: self }]) + } } #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, StableHash)] @@ -72,6 +78,11 @@ pub struct Type { pub kind: Kind, pub child: TypeTree, } +impl Type { + pub fn from_ty(offset: isize, other: &Type) -> Self { + Self { offset, size: other.size, kind: other.kind, child: other.child.clone() } + } +} impl Type { pub fn add_offset(self, add: isize) -> Self { diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 804547745c5d5..23cf01a84f7fd 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -1180,7 +1180,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { // 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); + crate::typetree::add_tt(self, memcpy, tt); } } diff --git a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs index ee17468ec0c03..10e152a7c8c63 100644 --- a/compiler/rustc_codegen_llvm/src/builder/autodiff.rs +++ b/compiler/rustc_codegen_llvm/src/builder/autodiff.rs @@ -143,7 +143,6 @@ pub(crate) fn adjust_activity_to_abi<'tcx>( // FIXME(ZuseZ4): This logic is a bit more complicated than it should be, can we simplify it // using iterators and peek()? fn match_args_from_caller_to_enzyme<'ll, 'tcx>( - cx: &SimpleCx<'ll>, builder: &mut Builder<'_, 'll, 'tcx>, width: u32, args: &mut Vec<&'ll Value>, @@ -157,6 +156,7 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>( // need to match those. // FIXME(ZuseZ4): This logic is a bit more complicated than it should be, can we simplify it // using iterators and peek()? + let cx = &builder.scx; let mut outer_pos: usize = 0; let mut activity_pos = 0; @@ -292,8 +292,7 @@ fn match_args_from_caller_to_enzyme<'ll, 'tcx>( // FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to // cover some assumptions of enzyme/autodiff, which could lead to UB otherwise. pub(crate) fn generate_enzyme_call<'ll, 'tcx>( - builder: &mut Builder<'_, 'll, 'tcx>, - cx: &SimpleCx<'ll>, + bx: &mut Builder<'_, 'll, 'tcx>, fn_to_diff: &'ll Value, outer_name: &str, ret_ty: &'ll Type, @@ -303,6 +302,7 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( dest_place: Option>, fnc_tree: FncTree, ) -> IntrinsicResult<'tcx, &'ll Value> { + let cx: &SimpleCx<'ll> = &bx.scx; // We have to pick the name depending on whether we want forward or reverse mode autodiff. let mut ad_name: String = match attrs.mode { DiffMode::Forward => "__enzyme_fwddiff", @@ -369,34 +369,27 @@ pub(crate) fn generate_enzyme_call<'ll, 'tcx>( args.push(cx.get_const_int(cx.type_i64(), attrs.width as u64)); } - match_args_from_caller_to_enzyme( - &cx, - builder, - attrs.width, - &mut args, - &attrs.input_activity, - fn_args, - ); + match_args_from_caller_to_enzyme(bx, attrs.width, &mut args, &attrs.input_activity, 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); + crate::typetree::add_tt(&bx, fn_to_diff, fnc_tree); } - let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None); + let call = bx.call(enzyme_ty, None, None, ad_fn, &args, None, None); - let fn_ret_ty = builder.cx.val_ty(call); - if fn_ret_ty == builder.cx.type_void() || fn_ret_ty == builder.cx.type_struct(&[], false) { + let fn_ret_ty = bx.cx.val_ty(call); + if fn_ret_ty == bx.cx.type_void() || fn_ret_ty == bx.cx.type_struct(&[], false) { // If we return void or an empty struct, then our caller (due to how we generated it) // does not expect a return value. As such, we have no pointer (or place) into which // we could store our value, and would store into an undef, which would cause UB. // As such, we just ignore the return value in those cases. IntrinsicResult::Operand(OperandValue::ZeroSized) } else if let Some(dest_place) = dest_place { - builder.store_to_place(call, dest_place); + bx.store_to_place(call, dest_place); IntrinsicResult::WroteIntoPlace } else { IntrinsicResult::Operand( - OperandRef::from_immediate_or_packed_pair(builder, call, dest_layout).val, + OperandRef::from_immediate_or_packed_pair(bx, call, dest_layout).val, ) } } diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 7bd604bdbbd76..bd7a4cfe38d48 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -225,7 +225,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { ) } sym::autodiff => { - return codegen_autodiff(self, tcx, instance, args, result_layout, result_place); + return codegen_autodiff(self, instance, args, result_layout, result_place); } sym::offload => { if tcx.sess.opts.unstable_opts.offload.is_empty() { @@ -1743,12 +1743,12 @@ fn codegen_retag_inner<'ll, 'tcx>( fn codegen_autodiff<'ll, 'tcx>( bx: &mut Builder<'_, 'll, 'tcx>, - tcx: TyCtxt<'tcx>, instance: ty::Instance<'tcx>, args: &[OperandRef<'tcx, &'ll Value>], result_layout: ty::layout::TyAndLayout<'tcx>, result_place: Option>, ) -> IntrinsicResult<'tcx, &'ll Value> { + let tcx = bx.tcx; if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) { let _ = tcx.dcx().emit_almost_fatal(AutoDiffWithoutEnable); } @@ -1815,7 +1815,6 @@ fn codegen_autodiff<'ll, 'tcx>( // Build body generate_enzyme_call( bx, - bx.cx, fn_to_diff, &diff_symbol, llret_ty, diff --git a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs index 195e050a9b651..68ec7870811d2 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs @@ -66,6 +66,7 @@ unsafe extern "C" { NameLen: libc::size_t, ) -> Option<&Value>; + pub(crate) safe fn LLVMRustIsCall(V: &Value) -> bool; } unsafe extern "C" { @@ -292,6 +293,11 @@ pub(crate) mod Enzyme_AD { unsafe { (self.EnzymeTypeTreeToString)(tree) } } + pub(crate) fn tree_to_cstr(&self, tree: *mut EnzymeTypeTree) -> &std::ffi::CStr { + let c_str = self.tree_to_string(tree); + unsafe { std::ffi::CStr::from_ptr(c_str) } + } + pub(crate) fn tree_to_string_free(&self, ch: *const c_char) { unsafe { (self.EnzymeTypeTreeToStringFree)(ch) } } diff --git a/compiler/rustc_codegen_llvm/src/llvm/mod.rs b/compiler/rustc_codegen_llvm/src/llvm/mod.rs index e9bc6ae0e80ef..a2d17e93b4996 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/mod.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/mod.rs @@ -76,6 +76,21 @@ pub(crate) fn CreateAttrStringValue<'ll>( ) } } +pub(crate) fn CreateAttrStringValueFromCStr<'ll>( + llcx: &'ll Context, + attr: &std::ffi::CStr, + value: &std::ffi::CStr, +) -> &'ll Attribute { + unsafe { + LLVMCreateStringAttribute( + llcx, + (*attr).as_ptr(), + (*attr).to_bytes().len() as c_uint, + (*value).as_ptr(), + (*value).to_bytes().len() as c_uint, + ) + } +} pub(crate) fn CreateAttrString<'ll>(llcx: &'ll Context, attr: &str) -> &'ll Attribute { unsafe { diff --git a/compiler/rustc_codegen_llvm/src/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs index 4f433f273c8cc..7c2e09227e46b 100644 --- a/compiler/rustc_codegen_llvm/src/typetree.rs +++ b/compiler/rustc_codegen_llvm/src/typetree.rs @@ -1,35 +1,47 @@ -use std::ffi::{CString, c_char, c_uint}; +use std::ffi::{CString, c_char}; -use rustc_ast::expand::typetree::{FncTree, TypeTree as RustTypeTree}; +use rustc_ast::expand::typetree::{FncTree, Kind, TypeTree as RustTypeTree}; use crate::attributes; +use crate::context::FullCx; use crate::llvm::{self, EnzymeWrapper, Value}; fn to_enzyme_typetree( - rust_typetree: RustTypeTree, + rust_typetree: &RustTypeTree, _data_layout: &str, llcx: &llvm::Context, -) -> llvm::TypeTree { +) -> (llvm::TypeTree, Vec) { let mut enzyme_tt = llvm::TypeTree::new(); - process_typetree_recursive(&mut enzyme_tt, &rust_typetree, &[], llcx); - enzyme_tt + let extra_ints = process_typetree_recursive(&mut enzyme_tt, &rust_typetree, &[], llcx); + + let mut int_vec = vec![]; + for _ in 0..extra_ints { + let mut int_tt = llvm::TypeTree::new(); + int_tt.insert(&[0], llvm::CConcreteType::DT_Integer, llcx); + int_vec.push(int_tt); + } + + (enzyme_tt, int_vec) } + fn process_typetree_recursive( enzyme_tt: &mut llvm::TypeTree, rust_typetree: &RustTypeTree, parent_indices: &[i64], llcx: &llvm::Context, -) { +) -> u32 { + let mut extra_ints = 0; 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, + Kind::Anything => llvm::CConcreteType::DT_Anything, + Kind::Integer => llvm::CConcreteType::DT_Integer, + Kind::Pointer => llvm::CConcreteType::DT_Pointer, + Kind::RustSlice => llvm::CConcreteType::DT_Pointer, + Kind::Half => llvm::CConcreteType::DT_Half, + Kind::Float => llvm::CConcreteType::DT_Float, + Kind::Double => llvm::CConcreteType::DT_Double, + Kind::F128 => llvm::CConcreteType::DT_FP128, + Kind::Unknown => llvm::CConcreteType::DT_Unknown, }; let mut indices = parent_indices.to_vec(); @@ -43,21 +55,28 @@ fn process_typetree_recursive( enzyme_tt.insert(&indices, concrete_type, llcx); - if rust_type.kind == rustc_ast::expand::typetree::Kind::Pointer + if matches!(rust_type.kind, Kind::RustSlice) { + // We lower slices to `ptr,int`, so add the int here. + extra_ints += 1; + } + + if matches!(rust_type.kind, Kind::Pointer | Kind::RustSlice) && !rust_type.child.0.is_empty() { process_typetree_recursive(enzyme_tt, &rust_type.child, &indices, llcx); } } + extra_ints +} + +// Describes all the locations in which we know how to apply an Enzyme TypeTree. +enum TTLocation { + Definition, + Callsite, } #[cfg_attr(not(feature = "llvm_enzyme"), allow(unused))] -pub(crate) fn add_tt<'ll>( - llmod: &'ll llvm::Module, - llcx: &'ll llvm::Context, - fn_def: &'ll Value, - tt: FncTree, -) { +pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: FncTree) { // TypeTree processing uses functions from Enzyme, which we might not have available if we did // not build this compiler with `llvm_enzyme`. This feature is not strictly necessary, but // skipping this function increases the chance that Enzyme fails to compile some code. @@ -66,6 +85,16 @@ pub(crate) fn add_tt<'ll>( #[cfg(not(feature = "llvm_enzyme"))] return; + let tcx = cx.tcx; + if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) { + return; + } + if tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::NoTT) { + return; + } + + let llmod = cx.llmod; + let llcx = cx.llcx; let inputs = tt.args; let ret_tt: RustTypeTree = tt.ret; @@ -77,41 +106,81 @@ pub(crate) fn add_tt<'ll>( let attr_name = "enzyme_type"; let c_attr_name = CString::new(attr_name).unwrap(); + let tt_location: TTLocation = + if llvm::LLVMRustIsCall(fn_def) { TTLocation::Callsite } else { TTLocation::Definition }; + + let mut offset = 0; for (i, input) in inputs.iter().enumerate() { - unsafe { - let enzyme_tt = to_enzyme_typetree(input.clone(), llvm_data_layout, llcx); + let (enzyme_tt, extra_ints) = to_enzyme_typetree(&input, llvm_data_layout, llcx); + + // This scope is just a visual reminder that we *must* drop the enzyme_wrapper before + // we drop any typetrees (mainly enzyme_tt and extra_ints). Drop calls can not accept + // arguments like an enzyme_wrapper, so the typetree drop impl has to call get_instance + // on the static enzyme instance, which is behind a Mutex. Therefore we'd deadlock if we + // hold the enzyme_wrapper while dropping the typetrees. + { let enzyme_wrapper = EnzymeWrapper::get_instance(); - let c_str = enzyme_wrapper.tree_to_string(enzyme_tt.inner); - let c_str = std::ffi::CStr::from_ptr(c_str); - - 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, - ); - - attributes::apply_to_llfn(fn_def, llvm::AttributePlace::Argument(i as u32), &[attr]); + let c_str = enzyme_wrapper.tree_to_cstr(enzyme_tt.inner); + + let attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str); + let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset); + // FIXME(autodiff): We currently know that this is correct for all the cases in which we + // call this function. But we should make it more robust for the future. + match tt_location { + TTLocation::Definition => { + attributes::apply_to_llfn(fn_def, arg_pos, &[attr]); + } + TTLocation::Callsite => { + attributes::apply_to_callsite(fn_def, arg_pos, &[attr]); + } + } enzyme_wrapper.tree_to_string_free(c_str.as_ptr()); + for v in &extra_ints { + offset += 1; + let c_str = enzyme_wrapper.tree_to_cstr(v.inner); + let int_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str); + let arg_pos = llvm::AttributePlace::Argument(i as u32 + offset); + match tt_location { + TTLocation::Definition => { + attributes::apply_to_llfn(fn_def, arg_pos, &[int_attr]); + } + TTLocation::Callsite => { + attributes::apply_to_callsite(fn_def, arg_pos, &[int_attr]); + } + } + enzyme_wrapper.tree_to_string_free(c_str.as_ptr()); + } + } + } + // We will only fail this if Rust types got lowered to LLVM in a way that we didn't predict. + // Error, so we can learn from our mistakes. + if matches!(tt_location, TTLocation::Definition) { + let expected = offset as usize + inputs.len(); + let actual = llvm::count_params(fn_def) as usize; + if expected != actual { + tcx.dcx().warn(format!( + "autodiff type-tree failure. We expected {expected} LLVM argument(s), \ + but the generated LLVM function has {actual} parameter(s)" + )); } } - unsafe { - let enzyme_tt = to_enzyme_typetree(ret_tt, llvm_data_layout, llcx); + // FIXME(autodiff): We should think more about what it means if a function returns a slice or + // other fat ptrs. + let (enzyme_tt, _extra_ints) = to_enzyme_typetree(&ret_tt, llvm_data_layout, llcx); + if ret_tt != RustTypeTree::new() { let enzyme_wrapper = EnzymeWrapper::get_instance(); - let c_str = enzyme_wrapper.tree_to_string(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, - ); - - attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]); + let c_str = enzyme_wrapper.tree_to_cstr(enzyme_tt.inner); + let ret_attr = llvm::CreateAttrStringValueFromCStr(llcx, &c_attr_name, &c_str); + let arg_pos = llvm::AttributePlace::ReturnValue; + match tt_location { + TTLocation::Definition => { + attributes::apply_to_llfn(fn_def, arg_pos, &[ret_attr]); + } + TTLocation::Callsite => { + attributes::apply_to_callsite(fn_def, arg_pos, &[ret_attr]); + } + } enzyme_wrapper.tree_to_string_free(c_str.as_ptr()); } } diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs index 39c2529152566..0d27ff90991b3 100644 --- a/compiler/rustc_codegen_ssa/src/traits/builder.rs +++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs @@ -2,10 +2,12 @@ use std::assert_matches; use std::ops::Deref; use rustc_abi::{Align, Scalar, Size, WrappingRange}; +use rustc_ast::expand::typetree::{FncTree, TypeTree}; use rustc_hir::attrs::AttributeKind; use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs; use rustc_middle::mir; use rustc_middle::ty::layout::{FnAbiOf, LayoutOf, TyAndLayout}; +use rustc_middle::ty::typetree::typetree_from_ty; use rustc_middle::ty::{AtomicOrdering, Instance, Ty}; use rustc_session::config::OptLevel; use rustc_span::Span; @@ -456,7 +458,7 @@ pub trait BuilderMethods<'a, 'tcx>: src_align: Align, size: Self::Value, flags: MemFlags, - tt: Option, + tt: Option, ); fn memmove( &mut self, @@ -517,6 +519,10 @@ pub trait BuilderMethods<'a, 'tcx>: let temp = self.load_operand(src.with_type(layout)); temp.val.store_with_flags(self, dst.with_type(layout), flags); } else if !layout.is_zst() { + let tt = typetree_from_ty(self.tcx(), layout.ty); + // We seem to pass all values to memcpy with one more indirection. + let tt = tt.add_indirection(); + let fnc_tree = FncTree { args: vec![tt.clone(), tt], ret: TypeTree::new() }; let bytes = self.const_usize(layout.size.bytes()); let bytes = if layout.peel_transparent_wrappers(self).ty.is_scalable_vector() { let vscale = self.vscale(self.type_i64()); @@ -524,7 +530,7 @@ pub trait BuilderMethods<'a, 'tcx>: } else { bytes }; - self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags, None); + self.memcpy(dst.llval, dst.align, src.llval, src.align, bytes, flags, Some(fnc_tree)); } } diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp index 3b8e6f6415365..f500041a12d8b 100644 --- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp +++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp @@ -160,6 +160,10 @@ extern "C" void LLVMRustPrintStatisticsJSON(RustStringRef OutBuf) { llvm::PrintStatisticsJSON(OS); } +extern "C" bool LLVMRustIsCall(LLVMValueRef V) { + return llvm::isa(llvm::unwrap(V)); +} + // Some of the functions here rely on LLVM modules that may not always be // available. As such, we only try to build it in the first place, if // llvm.offload is enabled. diff --git a/compiler/rustc_middle/src/ty/typetree.rs b/compiler/rustc_middle/src/ty/typetree.rs index 9e941bdb849ec..d4cda033a7e87 100644 --- a/compiler/rustc_middle/src/ty/typetree.rs +++ b/compiler/rustc_middle/src/ty/typetree.rs @@ -32,12 +32,19 @@ pub fn fnc_typetrees<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>) -> FncTree { // Create TypeTree for return type let ret = typetree_from_ty(tcx, sig.output()); - FncTree { args, ret } + let f = FncTree { args, ret }; + f } /// Generate a TypeTree for a specific type. /// Mainly a convenience wrapper around the actual implementation. pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree { + if !tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::Enable) { + return TypeTree::new(); + } + if tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::NoTT) { + return TypeTree::new(); + } let mut visited = Vec::new(); typetree_from_ty_impl_inner(tcx, ty, 0, &mut visited, false) } @@ -46,6 +53,40 @@ pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree { /// from pathological deeply nested types. Combined with cycle detection. const MAX_TYPETREE_DEPTH: usize = 6; +fn handle_indirection<'a>( + ty: Ty<'a>, + tcx: TyCtxt<'a>, + depth: usize, + visited: &mut Vec>, +) -> TypeTree { + let Some(inner_ty) = ty.builtin_deref(true) else { + bug!("incorrect autodiff typetree handling for type: {}", ty); + }; + // slices are represented as `&'{erased} mut [f32]` + // This reads as a reference to a slice of f32. + // So we'd end up with ptr->RustSlice->f32 without this extra handling + if inner_ty.is_slice() { + if let ty::Slice(element_ty) = inner_ty.kind() { + let element_tree = + typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false); + return TypeTree(vec![Type { + offset: -1, + size: tcx.data_layout.pointer_size().bytes_usize(), + kind: Kind::RustSlice, + child: element_tree, + }]); + } + } + + let child = typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true); + return TypeTree(vec![Type { + offset: -1, + size: tcx.data_layout.pointer_size().bytes_usize(), + kind: Kind::Pointer, + child, + }]); +} + /// Internal implementation with context about whether this is for a reference target. fn typetree_from_ty_impl_inner<'tcx>( tcx: TyCtxt<'tcx>, @@ -64,43 +105,12 @@ fn typetree_from_ty_impl_inner<'tcx>( } visited.push(ty); - 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), - _ => (Kind::Integer, 0), - } - } else { - (Kind::Integer, 0) - }; - - // Use offset 0 for scalars that are direct targets of references (like &f64) - // Use offset -1 for scalars used directly (like function return types) - let offset = if is_reference_target && !ty.is_array() { 0 } else { -1 }; - return TypeTree(vec![Type { offset, size, kind, child: TypeTree::new() }]); - } - - if ty.is_ref() || ty.is_raw_ptr() || ty.is_box() { - let Some(inner_ty) = ty.builtin_deref(true) else { - return TypeTree::new(); - }; - - let child = typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true); - 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() { + match ty.kind() { + // See handle_indirection for an explanation on why we don't handle it here. + ty::Slice(..) => bug!("incorrect autodiff typetree handling for slice: {}", ty), + ty::Ref(..) | ty::RawPtr(..) => handle_indirection(ty, tcx, depth, visited), + ty::Adt(def, _) if def.is_box() => handle_indirection(ty, tcx, depth, visited), + ty::Array(element_ty, len_const) => { let len = len_const.try_to_target_usize(tcx).unwrap_or(0); if len == 0 { return TypeTree::new(); @@ -109,65 +119,44 @@ fn typetree_from_ty_impl_inner<'tcx>( typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false); let mut types = Vec::new(); for elem_type in &element_tree.0 { - types.push(Type { - offset: -1, - size: elem_type.size, - kind: elem_type.kind, - child: elem_type.child.clone(), - }); + types.push(Type::from_ty(-1, elem_type)); } - return TypeTree(types); - } - } - - if ty.is_slice() { - if let ty::Slice(element_ty) = ty.kind() { - let element_tree = - typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false); - return element_tree; - } - } - - if let ty::Tuple(tuple_types) = ty.kind() { - if tuple_types.is_empty() { - return TypeTree::new(); + TypeTree(types) } + ty::Tuple(tuple_types) => { + if tuple_types.is_empty() { + return TypeTree::new(); + } - let mut types = Vec::new(); - let mut current_offset = 0; + let mut types = Vec::new(); + let mut current_offset = 0; - for tuple_ty in tuple_types.iter() { - let element_tree = - typetree_from_ty_impl_inner(tcx, tuple_ty, depth + 1, visited, false); + for tuple_ty in tuple_types.iter() { + let element_tree = + typetree_from_ty_impl_inner(tcx, tuple_ty, depth + 1, visited, false); - let element_layout = tcx - .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty)) - .ok() - .map(|layout| layout.size.bytes_usize()) - .unwrap_or(0); + let element_layout = tcx + .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty)) + .ok() + .map(|layout| layout.size.bytes_usize()) + .unwrap_or(0); - for elem_type in &element_tree.0 { - types.push(Type { - offset: if elem_type.offset == -1 { + for elem_type in &element_tree.0 { + let offset = if elem_type.offset == -1 { current_offset as isize } else { current_offset as isize + elem_type.offset - }, - size: elem_type.size, - kind: elem_type.kind, - child: elem_type.child.clone(), - }); + }; + types.push(Type::from_ty(offset, elem_type)); + } + + current_offset += element_layout; } - current_offset += element_layout; + TypeTree(types) } - - return TypeTree(types); - } - - if let ty::Adt(adt_def, args) = ty.kind() { - if adt_def.is_struct() { + ty::Adt(adt_def, args) if adt_def.is_struct() => { let struct_layout = tcx.layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(ty)); if let Ok(layout) = struct_layout { @@ -186,23 +175,37 @@ fn typetree_from_ty_impl_inner<'tcx>( let field_offset = layout.fields.offset(field_idx).bytes_usize(); for elem_type in &field_tree.0 { - types.push(Type { - offset: if elem_type.offset == -1 { - field_offset as isize - } else { - field_offset as isize + elem_type.offset - }, - size: elem_type.size, - kind: elem_type.kind, - child: elem_type.child.clone(), - }); + let offset = if elem_type.offset == -1 { + field_offset as isize + } else { + field_offset as isize + elem_type.offset + }; + types.push(Type::from_ty(offset, elem_type)); } } - return TypeTree(types); + TypeTree(types) + } else { + TypeTree::new() } } + ty::Char | ty::Bool | ty::Infer(ty::IntVar(_)) | ty::Int(_) | ty::Uint(_) => { + let kind = Kind::Integer; + let size = ty.primitive_size(tcx).bytes_usize(); + let offset = if is_reference_target { 0 } else { -1 }; + TypeTree(vec![Type { offset, size, kind, child: TypeTree::new() }]) + } + ty::Float(_) | ty::Infer(ty::FloatVar(_)) => { + let (enzyme_ty, size) = 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), + _ => bug!("Unexpected floating point type: {:?}", ty), + }; + let offset = if is_reference_target { 0 } else { -1 }; + TypeTree(vec![Type { offset, size, kind: enzyme_ty, child: TypeTree::new() }]) + } + _ => TypeTree::new(), } - - TypeTree::new() } diff --git a/tests/run-make/autodiff/type-trees/iter/rmake.rs b/tests/run-make/autodiff/type-trees/iter/rmake.rs new file mode 100644 index 0000000000000..4aca0771b7ee8 --- /dev/null +++ b/tests/run-make/autodiff/type-trees/iter/rmake.rs @@ -0,0 +1,25 @@ +//@ needs-enzyme +//@ ignore-cross-compile + +use run_make_support::{llvm_filecheck, rfs, rustc}; + +// This test passes in release mode. If we run it in Debug mode and don't lower any MIR info to +// LLVM TypeTrees, then it fails on deducing the type of a memcpy. If we lower info it still fails, +// but at a later location based on an extractvalue call. We will fix this in a future PR. + +fn main() { + rustc() + .input("window.rs") + .arg("-Zautodiff=Enable,NoTT") + .arg("-Clto=fat") + .run_fail() + .assert_stderr_contains("Enzyme: Cannot deduce type of copy"); + rustc() + .input("window.rs") + .arg("-Zautodiff=Enable") + .arg("-Clto=fat") + .emit("llvm-ir") + .run_fail() + .assert_stderr_contains("Enzyme: Cannot deduce type of extract"); + rustc().input("window.rs").arg("-Zautodiff=Enable,NoTT").arg("-Clto=fat").arg("-O").run(); +} diff --git a/tests/run-make/autodiff/type-trees/iter/window.rs b/tests/run-make/autodiff/type-trees/iter/window.rs new file mode 100644 index 0000000000000..7a6a6239fc23d --- /dev/null +++ b/tests/run-make/autodiff/type-trees/iter/window.rs @@ -0,0 +1,53 @@ +#![feature(autodiff)] + +use std::autodiff::autodiff_reverse; + +// This tests verifies that Enzyme can differentiate the iterator and window version of the for +// loops given below. Iterators (especially the windows use here) cause a lot of extra abstractions +// and indirections. Without extra typetree hints, Enzyme failed to differentiate them in debug +// mode. + +//@revisions: tt no_tt +//@[tt] compile-flags: -Z autodiff=Enable +//@[no_tt] compile-flags: -Z autodiff=Enable,NoTT +//@[no_tt] build-fail + +#[unsafe(no_mangle)] +#[inline(never)] +#[autodiff_reverse(f_rev, 2, Duplicated, Const, Duplicated)] +fn f(x: &[f64; 3], args: &[f64; 3], y: &mut [f64; 2]) { + y[0] = x.iter().map(|i| args[0] * i.powi(2)).sum(); + y[1] = x + .windows(2) + .map(|w| (args[1] - w[0]).powi(2) + args[2] * (w[1] - w[0].powi(2)).powi(2)) + .sum(); + // The iterators above are equivalent to the two following for loops. + // for i in 0..3 { + // y[0] += args[0] * x[i].powi(2); + // } + // for i in 0..2 { + // y[1] += (args[1] - x[i]).powi(2) + args[2] * (x[i + 1] - x[i].powi(2)).powi(2); + // } +} + +// Not generally recommended, but since we rewrite llvm-ir, it should be good enough. +fn assert_abs_diff_eq(x: &[f64; N], y: &[f64; N]) { + for i in 0..N { + assert_eq!(x[i], y[i]); + } +} + +fn main() { + let x = [3.0, 5.0, 7.0]; + let args = [2.0, 1.0, 100.0]; + + let mut vjp = ([0.0; 3], [0.0; 3]); + let mut y = [0.0; 2]; + let mut dy = ([1.0, 0.0], [0.0, 1.0]); + + f_rev(&x, &mut vjp.0, &mut vjp.1, &args, &mut y, &mut dy.0, &mut dy.1); + + assert_abs_diff_eq::<2>(&y, &[166.0, 34020.0]); + assert_abs_diff_eq::<3>(&vjp.0, &[12.0, 20.0, 28.0]); + assert_abs_diff_eq::<3>(&vjp.1, &[4804.0, 35208.0, -3600.0]); +}