diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs index 8f192a5a73b63..c458968c15370 100644 --- a/compiler/rustc_codegen_llvm/src/asm.rs +++ b/compiler/rustc_codegen_llvm/src/asm.rs @@ -365,7 +365,7 @@ impl<'ll, 'tcx> AsmBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { let value = if output_types.len() == 1 { result } else { - self.extract_value(result, op_idx[&idx] as u64) + self.extract_value(result, op_idx[&idx] as u64, None) }; let value = llvm_fixup_output(self, value, reg.reg_class(), &place.layout, instance); diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs index 23cf01a84f7fd..384b1c6e3e945 100644 --- a/compiler/rustc_codegen_llvm/src/builder.rs +++ b/compiler/rustc_codegen_llvm/src/builder.rs @@ -2,7 +2,7 @@ use std::borrow::{Borrow, Cow}; use std::iter; use std::ops::Deref; -use rustc_ast::expand::typetree::FncTree; +use rustc_ast::expand::typetree::{FncTree, TypeTree}; pub(crate) mod autodiff; pub(crate) mod gpu_offload; @@ -636,7 +636,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { let name = format!("llvm.{}{oop_str}.with.overflow", if signed { 's' } else { 'u' }); let res = self.call_intrinsic(name, &[self.type_ix(width)], &[lhs, rhs]); - (self.extract_value(res, 0), self.extract_value(res, 1)) + (self.extract_value(res, 0, None), self.extract_value(res, 1, None)) } fn from_immediate(&mut self, val: Self::Value) -> Self::Value { @@ -1262,9 +1262,19 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { } } - fn extract_value(&mut self, agg_val: &'ll Value, idx: u64) -> &'ll Value { + fn extract_value(&mut self, agg_val: &'ll Value, idx: u64, tt: Option) -> &'ll Value { assert_eq!(idx as c_uint as u64, idx); - unsafe { llvm::LLVMBuildExtractValue(self.llbuilder, agg_val, idx as c_uint, UNNAMED) } + let extract = + unsafe { llvm::LLVMBuildExtractValue(self.llbuilder, agg_val, idx as c_uint, UNNAMED) }; + if let Some(tt) = tt { + let fnc_tree = FncTree { args: vec![], ret: tt }; + crate::typetree::add_tt( + self, + extract, + fnc_tree, + ); + } + extract } fn insert_value(&mut self, agg_val: &'ll Value, elt: &'ll Value, idx: u64) -> &'ll Value { @@ -1284,7 +1294,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { unsafe { llvm::LLVMSetCleanup(landing_pad, llvm::TRUE); } - (self.extract_value(landing_pad, 0), self.extract_value(landing_pad, 1)) + (self.extract_value(landing_pad, 0, None), self.extract_value(landing_pad, 1, None)) } fn filter_landing_pad(&mut self, pers_fn: &'ll Value) { @@ -1385,8 +1395,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { llvm::FALSE, // SingleThreaded ); llvm::LLVMSetWeak(value, weak.to_llvm_bool()); - let val = self.extract_value(value, 0); - let success = self.extract_value(value, 1); + let val = self.extract_value(value, 0, None); + let success = self.extract_value(value, 1, None); (val, success) } } diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 0b009321802cf..39e3ac8eb484f 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs @@ -162,9 +162,9 @@ impl<'ll> OffloadKernelDims<'ll> { builder: &mut Builder<'_, 'll, 'tcx>, arr: &'ll Value, ) -> &'ll Value { - let x = builder.extract_value(arr, 0); - let y = builder.extract_value(arr, 1); - let z = builder.extract_value(arr, 2); + let x = builder.extract_value(arr, 0, None); + let y = builder.extract_value(arr, 1, None); + let z = builder.extract_value(arr, 2, None); let xy = builder.mul(x, y); builder.mul(xy, z) diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index bd7a4cfe38d48..72e3eeaf459b9 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -791,6 +791,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { self.extract_value( args[0].immediate(), fn_args.const_at(2).to_leaf().to_i32() as u64, + None, ) } @@ -1074,7 +1075,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> { &[], &[llvtable, vtable_byte_offset, typeid], ); - self.extract_value(type_checked_load, 0) + self.extract_value(type_checked_load, 0, None) } fn va_start(&mut self, va_list: &'ll Value) { @@ -1186,7 +1187,7 @@ fn autocast<'ll>( iter::zip(bx.struct_element_types(src_ty), bx.struct_element_types(dest_ty)) .enumerate() { - let elt = bx.extract_value(val, idx as u64); + let elt = bx.extract_value(val, idx as u64, None); let casted_elt = autocast(bx, elt, src_element_ty, dest_element_ty); ret = bx.insert_value(ret, casted_elt, idx as u64); } @@ -1647,7 +1648,7 @@ fn codegen_gnu_try<'ll, 'tcx>( let vals = bx.landing_pad(lpad_ty, bx.eh_personality(), 1); let tydesc = bx.const_null(bx.type_ptr()); bx.add_clause(vals, tydesc); - let ptr = bx.extract_value(vals, 0); + let ptr = bx.extract_value(vals, 0, None); let catch_ty = bx.type_func(&[bx.type_ptr(), bx.type_ptr()], bx.type_void()); bx.call(catch_ty, None, None, catch_func, &[data, ptr], None, None); bx.ret(bx.const_bool(true)); @@ -1937,8 +1938,9 @@ fn get_args_from_tuple<'ll, 'tcx>( let field = tuple_place.project_field(bx, tuple_index); let llvm_ty = field.layout.llvm_type(bx.cx); let pair_val = bx.load(llvm_ty, field.val.llval, field.val.align); - result.push(bx.extract_value(pair_val, 0)); - result.push(bx.extract_value(pair_val, 1)); + + result.push(bx.extract_value(pair_val, 0, None)); + result.push(bx.extract_value(pair_val, 1, None)); tuple_index += 1; } PassMode::Indirect { .. } => { diff --git a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs index 68ec7870811d2..8e90b6ac84ffb 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs @@ -67,6 +67,8 @@ unsafe extern "C" { ) -> Option<&Value>; pub(crate) safe fn LLVMRustIsCall(V: &Value) -> bool; + pub(crate) safe fn LLVMRustSupportsEnzymeMD(V: &Value) -> bool; + pub(crate) fn LLVMRustSetEnzymeTypeMD(v: &Value, md: &Value); } unsafe extern "C" { @@ -97,7 +99,7 @@ pub(crate) mod Enzyme_AD { use rustc_session::filesearch; use super::{CConcreteType, CTypeTreeRef, Context}; - use crate::llvm::{EnzymeTypeTree, LLVMRustVersionMajor}; + use crate::llvm::{EnzymeTypeTree, LLVMRustVersionMajor, Value}; type EnzymeSetCLBoolFn = unsafe extern "C" fn(*mut c_void, u8); type EnzymeSetCLStringFn = unsafe extern "C" fn(*mut c_void, *const c_char); @@ -115,6 +117,7 @@ pub(crate) mod Enzyme_AD { unsafe extern "C" fn(CTypeTreeRef, *const i64, usize, CConcreteType, &Context); type EnzymeTypeTreeToStringFn = unsafe extern "C" fn(CTypeTreeRef) -> *const c_char; type EnzymeTypeTreeToStringFreeFn = unsafe extern "C" fn(*const c_char); + type EnzymeTypeTreeToMDFn = unsafe extern "C" fn(CTypeTreeRef, &Context) -> Option<&Value>; #[allow(non_snake_case)] pub(crate) struct EnzymeWrapper { @@ -129,6 +132,7 @@ pub(crate) mod Enzyme_AD { EnzymeTypeTreeInsertEq: EnzymeTypeTreeInsertEqFn, EnzymeTypeTreeToString: EnzymeTypeTreeToStringFn, EnzymeTypeTreeToStringFree: EnzymeTypeTreeToStringFreeFn, + EnzymeTypeTreeToMD: EnzymeTypeTreeToMDFn, EnzymePrintPerf: *mut c_void, EnzymePrintActivity: *mut c_void, @@ -302,6 +306,14 @@ pub(crate) mod Enzyme_AD { unsafe { (self.EnzymeTypeTreeToStringFree)(ch) } } + pub(crate) fn tree_to_md<'a>( + &'a self, + tree: *mut EnzymeTypeTree, + ctx: &'a Context, + ) -> Option<&'a Value> { + unsafe { (self.EnzymeTypeTreeToMD)(tree, ctx) } + } + pub(crate) fn get_max_type_depth(&self) -> usize { unsafe { std::ptr::read::(self.EnzymeMaxTypeDepth as *const u32) as usize } } @@ -385,6 +397,7 @@ pub(crate) mod Enzyme_AD { EnzymeTypeTreeInsertEq: EnzymeTypeTreeInsertEqFn, EnzymeTypeTreeToString: EnzymeTypeTreeToStringFn, EnzymeTypeTreeToStringFree: EnzymeTypeTreeToStringFreeFn, + EnzymeTypeTreeToMD: EnzymeTypeTreeToMDFn, EnzymeSetCLBool: EnzymeSetCLBoolFn, EnzymeSetCLString: EnzymeSetCLStringFn, ); @@ -416,6 +429,7 @@ pub(crate) mod Enzyme_AD { EnzymeTypeTreeInsertEq, EnzymeTypeTreeToString, EnzymeTypeTreeToStringFree, + EnzymeTypeTreeToMD, EnzymePrintPerf, EnzymePrintActivity, EnzymePrintType, diff --git a/compiler/rustc_codegen_llvm/src/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs index 7c2e09227e46b..e2c2496f4e535 100644 --- a/compiler/rustc_codegen_llvm/src/typetree.rs +++ b/compiler/rustc_codegen_llvm/src/typetree.rs @@ -73,6 +73,7 @@ fn process_typetree_recursive( enum TTLocation { Definition, Callsite, + Intrinsic, } #[cfg_attr(not(feature = "llvm_enzyme"), allow(unused))] @@ -106,14 +107,19 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: 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 tt_location: TTLocation = if llvm::LLVMRustIsCall(fn_def) { + TTLocation::Callsite + } else if llvm::LLVMRustSupportsEnzymeMD(fn_def) { + TTLocation::Intrinsic + } else { + TTLocation::Definition + }; let mut offset = 0; for (i, input) in inputs.iter().enumerate() { 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 + // This scope is a simple solution since 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 @@ -133,6 +139,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: TTLocation::Callsite => { attributes::apply_to_callsite(fn_def, arg_pos, &[attr]); } + TTLocation::Intrinsic => { + let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx); + unsafe { + llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap()); + } + } } enzyme_wrapper.tree_to_string_free(c_str.as_ptr()); for v in &extra_ints { @@ -141,6 +153,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: 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::Intrinsic => { + let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx); + unsafe { + llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap()); + } + } TTLocation::Definition => { attributes::apply_to_llfn(fn_def, arg_pos, &[int_attr]); } @@ -152,8 +170,8 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: } } } - // 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. + // We will fail here if Rust types got lowered to LLVM in a way that we didn't predict. + // We 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; @@ -180,6 +198,12 @@ pub(crate) fn add_tt<'tcx, 'll>(cx: &FullCx<'ll, 'tcx>, fn_def: &'ll Value, tt: TTLocation::Callsite => { attributes::apply_to_callsite(fn_def, arg_pos, &[ret_attr]); } + TTLocation::Intrinsic => { + let md = enzyme_wrapper.tree_to_md(enzyme_tt.inner, llcx); + unsafe { + llvm::LLVMRustSetEnzymeTypeMD(fn_def, md.unwrap()); + } + } } enzyme_wrapper.tree_to_string_free(c_str.as_ptr()); } diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs index fad93de43b272..5ee49d12b6479 100644 --- a/compiler/rustc_codegen_ssa/src/mir/block.rs +++ b/compiler/rustc_codegen_ssa/src/mir/block.rs @@ -2360,8 +2360,8 @@ pub fn store_cast<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>( if let Some(offset_from_start) = cast.rest_offset { assert_eq!(cast.prefix.len(), 1); assert_eq!(cast.rest.unit.size, cast.rest.total); - let first = bx.extract_value(value, 0); - let second = bx.extract_value(value, 1); + let first = bx.extract_value(value, 0, None); + let second = bx.extract_value(value, 1, None); bx.store(first, ptr, align); let second_ptr = bx.inbounds_ptradd(ptr, bx.const_usize(offset_from_start.bytes())); bx.store(second, second_ptr, align.restrict_for_offset(offset_from_start)); diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs index 0246da3f27829..e94e5daf399f0 100644 --- a/compiler/rustc_codegen_ssa/src/mir/operand.rs +++ b/compiler/rustc_codegen_ssa/src/mir/operand.rs @@ -5,21 +5,60 @@ use rustc_abi as abi; use rustc_abi::{ Align, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, TagEncoding, VariantIdx, Variants, }; +use rustc_ast::expand::typetree::TypeTree; use rustc_hir::LangItem; use rustc_middle::mir::interpret::{Pointer, Scalar, alloc_range}; use rustc_middle::mir::{self, ConstValue}; use rustc_middle::ty::layout::{LayoutOf, TyAndLayout}; +use rustc_middle::ty::typetree::typetree_from_ty; use rustc_middle::ty::{self, Ty}; use rustc_middle::{bug, span_bug}; use rustc_session::config::{AnnotateMoves, DebugInfo, OptLevel}; +use rustc_span::sym; use tracing::{debug, instrument}; use super::place::{PlaceRef, PlaceValue}; use super::rvalue::transmute_scalar; use super::{FunctionCx, LocalRef}; -use crate::MemFlags; use crate::common::IntPredicate; use crate::traits::*; +use crate::{MemFlags, TyCtxt}; + +fn option_ptr_like_scalar_pair_tts<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option { + let ty::Adt(def, args) = ty.kind() else { + return None; + }; + + if !tcx.is_lang_item(def.did(), LangItem::Option) { + return None; + } + + let inner = args.type_at(0); + if !(inner.is_ref() || inner.is_box() || nonnull_inner_ty(tcx, inner).is_some()) { + return None; + } + + let tt = typetree_from_ty(tcx, inner); + //let some_layout = layout.for_variant(bx.cx(), VariantIdx::from_u32(1)); + //let payload_layout = some_layout.field(bx.cx(), 0); + // this will be a slice + //let payload_ty = payload_layout.ty; + //let tt = rustc_middle::ty::typetree_from_ty(bx.tcx(), field0_ty.unwrap()); + if tt == TypeTree::new() { + return None; + } + Some(tt) +} + +fn nonnull_inner_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> Option> { + if let ty::Adt(def, args) = ty.kind() + && tcx.is_diagnostic_item(sym::NonNull, def.did()) + { + return Some(args.type_at(0)); + } + + None +} /// The representation of a Rust value. The enum variant is in fact /// uniquely determined by the value's type, but is kept as a @@ -349,8 +388,10 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> { debug!("Operand::from_immediate_or_packed_pair: unpacking {:?} @ {:?}", llval, layout); // Deconstruct the immediate aggregate. - let a_llval = bx.extract_value(llval, 0); - let b_llval = bx.extract_value(llval, 1); + let f1 = option_ptr_like_scalar_pair_tts(bx.tcx(), layout.ty); + let f2 = if f1.is_none() { None } else { Some(TypeTree::int(8)) }; + let a_llval = bx.extract_value(llval, 0, f1); + let b_llval = bx.extract_value(llval, 1, f2); OperandValue::Pair(a_llval, b_llval) } else { OperandValue::Immediate(llval) diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs index 2549d0a039109..5cea2a81882a1 100644 --- a/compiler/rustc_codegen_ssa/src/mir/statement.rs +++ b/compiler/rustc_codegen_ssa/src/mir/statement.rs @@ -2,6 +2,10 @@ use rustc_middle::mir::{self, NonDivergingIntrinsic, StmtDebugInfo}; use rustc_middle::{bug, span_bug, ty}; use tracing::instrument; +use rustc_ast::expand::typetree::{TypeTree, FncTree}; +use rustc_middle::ty::typetree::typetree_from_ty; + + use super::{FunctionCx, LocalRef}; use crate::mir::retag; use crate::traits::*; @@ -111,7 +115,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> { let dst = dst_val.immediate(); let src = src_val.immediate(); - bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), None); + let tt = typetree_from_ty(bx.tcx(), pointee).add_indirection(); + dbg!(&tt); + let fnc_tree = FncTree { + args: vec![tt.clone(), tt], + ret: TypeTree::new(), + }; + + bx.memcpy(dst, align, src, align, bytes, crate::MemFlags::empty(), Some(fnc_tree)); } mir::StatementKind::FakeRead(..) | mir::StatementKind::AscribeUserType(..) diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs index 0d27ff90991b3..19fd5da3cd004 100644 --- a/compiler/rustc_codegen_ssa/src/traits/builder.rs +++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs @@ -568,7 +568,12 @@ pub trait BuilderMethods<'a, 'tcx>: fn va_arg(&mut self, list: Self::Value, ty: Self::Type) -> Self::Value; fn extract_element(&mut self, vec: Self::Value, idx: Self::Value) -> Self::Value; fn vector_splat(&mut self, num_elts: usize, elt: Self::Value) -> Self::Value; - fn extract_value(&mut self, agg_val: Self::Value, idx: u64) -> Self::Value; + fn extract_value( + &mut self, + agg_val: Self::Value, + idx: u64, + tt: Option, + ) -> Self::Value; fn insert_value(&mut self, agg_val: Self::Value, elt: Self::Value, idx: u64) -> Self::Value; fn set_personality_fn(&mut self, personality: Self::Function); diff --git a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp index f500041a12d8b..0574a37997b09 100644 --- a/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp +++ b/compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp @@ -160,6 +160,25 @@ extern "C" void LLVMRustPrintStatisticsJSON(RustStringRef OutBuf) { llvm::PrintStatisticsJSON(OS); } +extern "C" bool LLVMRustSupportsEnzymeMD(LLVMValueRef V) { + auto *I = llvm::dyn_cast(llvm::unwrap(V)); + return I && llvm::isa(I); +} + +extern "C" void LLVMRustSetEnzymeTypeMD(LLVMValueRef V, LLVMValueRef MDV) { + llvm::errs() << "setting MD" << "\n"; + auto *I = llvm::dyn_cast(llvm::unwrap(V)); + assert(I && "expected instruction for !enzyme_type metadata"); + + auto *MAV = llvm::dyn_cast(llvm::unwrap(MDV)); + assert(MAV && "expected MetadataAsValue"); + + auto *MD = llvm::dyn_cast(MAV->getMetadata()); + assert(MD && "expected MDNode"); + + I->setMetadata("enzyme_type", MD); +} + extern "C" bool LLVMRustIsCall(LLVMValueRef V) { return llvm::isa(llvm::unwrap(V)); } diff --git a/tests/run-make/autodiff/type-trees/iter/rmake.rs b/tests/run-make/autodiff/type-trees/iter/rmake.rs index 4aca0771b7ee8..b25134f4a25ae 100644 --- a/tests/run-make/autodiff/type-trees/iter/rmake.rs +++ b/tests/run-make/autodiff/type-trees/iter/rmake.rs @@ -14,12 +14,7 @@ fn main() { .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").arg("-Clto=fat").emit("llvm-ir").run(); rustc().input("window.rs").arg("-Zautodiff=Enable,NoTT").arg("-Clto=fat").arg("-O").run(); + llvm_filecheck().patterns("window.check").stdin_buf(rfs::read("window.ll")).run(); } diff --git a/tests/run-make/autodiff/type-trees/iter/window.check b/tests/run-make/autodiff/type-trees/iter/window.check new file mode 100644 index 0000000000000..340fd4d6f1b9c --- /dev/null +++ b/tests/run-make/autodiff/type-trees/iter/window.check @@ -0,0 +1,24 @@ +; Check that array TypeTree metadata is correctly generated for the iter/window case +; We check all three relevant cases. The metadata generated for extractvalue, +; the metadata for the memcpy call, +; and the metadata for the function definition. +; At the time of writing, at least the first two where necessary to make compilation succeed in debug mode. + +CHECK-LABEL: ; as core::iter::traits::iterator::Iterator>::fold +CHECK: extractvalue { ptr, i64 } %[[IN:[0-9]+]], 0, !enzyme_type ![[DATA_TT:[0-9]+]] +CHECK: extractvalue { ptr, i64 } %[[IN]], 1, !enzyme_type ![[LEN_TT:[0-9]+]] + +CHECK-LABEL: ; window::main +CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 "enzyme_type"="{[0]:Pointer, [0,-1]:Float@double}" %vjp, ptr align 8 "enzyme_type"="{[0]:Pointer, [0,-1]:Float@double}" {{.*}}, i64 24, i1 false) + + +CHECK: define {{.*}}void @f(ptr align 8 "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %x, ptr align 8 "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %0, ptr writeonly align 8 captures(none) "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@double}" %y) + + +CHECK-DAG: ![[DATA_TT]] = !{!"Unknown", i32 -1, ![[PTR_TT:[0-9]+]]} +CHECK-DAG: ![[PTR_TT]] = !{!"Pointer", i32 -1, ![[FLOAT_TT:[0-9]+]]} +CHECK-DAG: ![[FLOAT_TT]] = !{!"Float@double"} + +CHECK-DAG: ![[LEN_TT]] = !{!"Unknown", i32 0, ![[INT_TT:[0-9]+]], i32 1, ![[INT_TT]], i32 2, ![[INT_TT]], i32 3, ![[INT_TT]], i32 4, ![[INT_TT]], i32 5, ![[INT_TT]], i32 6, ![[INT_TT]], i32 7, ![[INT_TT]]} +CHECK-DAG: ![[INT_TT]] = !{!"Integer"} + diff --git a/tests/run-make/autodiff/type-trees/iter2/asdf.rs b/tests/run-make/autodiff/type-trees/iter2/asdf.rs new file mode 100644 index 0000000000000..e0a0f413d1e2c --- /dev/null +++ b/tests/run-make/autodiff/type-trees/iter2/asdf.rs @@ -0,0 +1,63 @@ +#![feature(autodiff)] +use std::autodiff::*; +use std::f64::consts::FRAC_PI_6; +use ndarray::arr1; + +#[autodiff_forward(da_args, Dual, Dual)] +pub fn a_f64_args(args: &[f64]) -> f64 { + let temperature = args[0]; + let volume = args[1]; + let moles = &args[2..]; + + // Now this crashes as well ... + let m = arr1(&[2.001829]); + let sigma = arr1(&[3.618353]); + let epsilon_k = arr1(&[208.1101]); + + // ... as does this ... + //let m = vec![2.001829]; + //let sigma = vec![3.618353]; + //let epsilon_k = vec![208.1101]; + + // but this works + // let m = &[2.001829]; + // let sigma = &[3.618353]; + // let epsilon_k = &[208.1101]; + + let n = moles.len(); + let t_inv = temperature.recip(); + let diameter: Vec = (0..n) + .map(|i| -((t_inv * -3.0 * epsilon_k[i]).exp() * 0.12 - 1.0) * sigma[i]) + .collect(); + + //let partial_density: Vec = moles.iter().cloned().map(|n| n / volume).collect(); // <- works + let partial_density: Vec = moles.iter().map(|&n| n / volume).collect(); // <- crashes + let density: f64 = partial_density.iter().cloned().sum(); + let total_moles: f64 = moles.iter().cloned().sum(); + let x: Vec = moles.iter().cloned().map(|n| n / total_moles).collect(); + + let mut zeta = [0.0; 4]; + for i in 0..diameter.len() { + for (z, &k) in zeta.iter_mut().zip([0, 1, 2, 3].iter()) { + *z += x[i] * diameter[i].powi(k) * (m[i] * FRAC_PI_6); + } + } + let zeta_23 = zeta[2] / zeta[3]; + + zeta.iter_mut().for_each(|z| *z *= density); + let frac_1mz3 = -(zeta[3] - 1.0).recip(); + let a = volume / std::f64::consts::FRAC_PI_6 + * (zeta[1] * zeta[2] * frac_1mz3 * 3.0 + + zeta[2].powi(2) * frac_1mz3.powi(2) * zeta_23 + + (zeta[2] * zeta_23.powi(2) - zeta[0]) * (zeta[3] * (-1.0)).ln_1p()); + a +} + +fn main() { + let t = 250.0; + let v = 1000.0; + + let seed = &[1.0, 0.0, 0.0]; // first entry is temperature + let da_dt_enz_args = da_args(&[t, v, 1.0], seed); + dbg!(&da_dt_enz_args); +}