diff options
Diffstat (limited to 'compiler/rustc_codegen_llvm')
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/abi.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/back/lto.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/builder.rs | 15 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/builder/autodiff.rs | 6 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/debuginfo/mod.rs | 81 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/debuginfo/utils.rs | 2 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/intrinsic.rs | 4 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/lib.rs | 1 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs | 206 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/llvm/ffi.rs | 97 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/typetree.rs | 122 | ||||
| -rw-r--r-- | compiler/rustc_codegen_llvm/src/va_arg.rs | 1 |
13 files changed, 457 insertions, 83 deletions
diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/compiler/rustc_codegen_llvm/src/abi.rs index 861227f7c2a..1703cab942b 100644 --- a/compiler/rustc_codegen_llvm/src/abi.rs +++ b/compiler/rustc_codegen_llvm/src/abi.rs @@ -246,6 +246,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 78107d95e5a..5ac3a87c158 100644 --- a/compiler/rustc_codegen_llvm/src/back/lto.rs +++ b/compiler/rustc_codegen_llvm/src/back/lto.rs @@ -563,6 +563,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 a4dc4eb532f..5271d0b4bb8 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; @@ -1107,11 +1108,12 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> { src_align: Align, size: &'ll Value, flags: MemFlags, + tt: Option<FncTree>, ) { 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, @@ -1120,7 +1122,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 b66e3dfdeec..c3485f56391 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::{Instance, PseudoCanonicalInput, TyCtxt, TypingEnv}; @@ -294,6 +295,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 { @@ -370,6 +372,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/debuginfo/metadata.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs index 4ba72cd61a0..4b74c04ed7a 100644 --- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs +++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs @@ -117,7 +117,7 @@ fn build_fixed_size_array_di_node<'ll, 'tcx>( .try_to_target_usize(cx.tcx) .expect("expected monomorphic const in codegen") as c_longlong; - let subrange = unsafe { llvm::LLVMRustDIBuilderGetOrCreateSubrange(DIB(cx), 0, upper_bound) }; + let subrange = unsafe { llvm::LLVMDIBuilderGetOrCreateSubrange(DIB(cx), 0, upper_bound) }; let subscripts = &[subrange]; let di_node = unsafe { diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs index 126082aa3aa..af64e4ebed0 100644 --- a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs +++ b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs @@ -52,15 +52,6 @@ mod utils; use self::create_scope_map::compute_mir_scopes; pub(crate) use self::metadata::build_global_var_di_node; -// FIXME(Zalathar): These `DW_TAG_*` constants are fake values that were -// removed from LLVM in 2015, and are only used by our own `RustWrapper.cpp` -// to decide which C++ API to call. Instead, we should just have two separate -// FFI functions and choose the correct one on the Rust side. -#[allow(non_upper_case_globals)] -const DW_TAG_auto_variable: c_uint = 0x100; -#[allow(non_upper_case_globals)] -const DW_TAG_arg_variable: c_uint = 0x101; - /// A context object for maintaining all state needed by the debuginfo module. pub(crate) struct CodegenUnitDebugContext<'ll, 'tcx> { llmod: &'ll llvm::Module, @@ -174,35 +165,38 @@ impl<'ll> DebugInfoBuilderMethods for Builder<'_, 'll, '_> { if direct_offset.bytes() > 0 { addr_ops.push(DW_OP_plus_uconst); - addr_ops.push(direct_offset.bytes() as u64); + addr_ops.push(direct_offset.bytes()); } for &offset in indirect_offsets { addr_ops.push(DW_OP_deref); if offset.bytes() > 0 { addr_ops.push(DW_OP_plus_uconst); - addr_ops.push(offset.bytes() as u64); + addr_ops.push(offset.bytes()); } } if let Some(fragment) = fragment { // `DW_OP_LLVM_fragment` takes as arguments the fragment's // offset and size, both of them in bits. addr_ops.push(DW_OP_LLVM_fragment); - addr_ops.push(fragment.start.bits() as u64); - addr_ops.push((fragment.end - fragment.start).bits() as u64); + addr_ops.push(fragment.start.bits()); + addr_ops.push((fragment.end - fragment.start).bits()); } + let di_builder = DIB(self.cx()); + let addr_expr = unsafe { + llvm::LLVMDIBuilderCreateExpression(di_builder, addr_ops.as_ptr(), addr_ops.len()) + }; unsafe { // FIXME(eddyb) replace `llvm.dbg.declare` with `llvm.dbg.addr`. - llvm::LLVMRustDIBuilderInsertDeclareAtEnd( - DIB(self.cx()), + llvm::LLVMDIBuilderInsertDeclareRecordAtEnd( + di_builder, variable_alloca, dbg_var, - addr_ops.as_ptr(), - addr_ops.len() as c_uint, + addr_expr, dbg_loc, self.llbb(), - ); - } + ) + }; } fn set_dbg_loc(&mut self, dbg_loc: &'ll DILocation) { @@ -630,28 +624,39 @@ impl<'ll, 'tcx> DebugInfoCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> { let type_metadata = spanned_type_di_node(self, variable_type, span); - let (argument_index, dwarf_tag) = match variable_kind { - ArgumentVariable(index) => (index as c_uint, DW_TAG_arg_variable), - LocalVariable => (0, DW_TAG_auto_variable), - }; let align = self.align_of(variable_type); let name = variable_name.as_str(); - unsafe { - llvm::LLVMRustDIBuilderCreateVariable( - DIB(self), - dwarf_tag, - scope_metadata, - name.as_c_char_ptr(), - name.len(), - file_metadata, - loc.line, - type_metadata, - true, - DIFlags::FlagZero, - argument_index, - align.bits() as u32, - ) + + match variable_kind { + ArgumentVariable(arg_index) => unsafe { + llvm::LLVMDIBuilderCreateParameterVariable( + DIB(self), + scope_metadata, + name.as_ptr(), + name.len(), + arg_index as c_uint, + file_metadata, + loc.line, + type_metadata, + llvm::Bool::TRUE, // (preserve descriptor during optimizations) + DIFlags::FlagZero, + ) + }, + LocalVariable => unsafe { + llvm::LLVMDIBuilderCreateAutoVariable( + DIB(self), + scope_metadata, + name.as_ptr(), + name.len(), + file_metadata, + loc.line, + type_metadata, + llvm::Bool::TRUE, // (preserve descriptor during optimizations) + DIFlags::FlagZero, + align.bits() as u32, + ) + }, } } } diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/utils.rs b/compiler/rustc_codegen_llvm/src/debuginfo/utils.rs index cc1d504b430..7e1e49310f6 100644 --- a/compiler/rustc_codegen_llvm/src/debuginfo/utils.rs +++ b/compiler/rustc_codegen_llvm/src/debuginfo/utils.rs @@ -28,7 +28,7 @@ pub(crate) fn create_DIArray<'ll>( builder: &DIBuilder<'ll>, arr: &[Option<&'ll DIDescriptor>], ) -> &'ll DIArray { - unsafe { llvm::LLVMRustDIBuilderGetOrCreateArray(builder, arr.as_ptr(), arr.len() as u32) } + unsafe { llvm::LLVMDIBuilderGetOrCreateArray(builder, arr.as_ptr(), arr.len()) } } #[inline] diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 50398a32142..013108d1286 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/compiler/rustc_codegen_llvm/src/intrinsic.rs @@ -1212,6 +1212,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, @@ -1222,6 +1225,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 13bdb7cb1a2..b9584a9117f 100644 --- a/compiler/rustc_codegen_llvm/src/lib.rs +++ b/compiler/rustc_codegen_llvm/src/lib.rs @@ -68,6 +68,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 695435eb6da..e63043b2122 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,40 @@ 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 EnzymeTypeTreeInsertEq( + CTT: CTypeTreeRef, + indices: *const i64, + len: usize, + ct: CConcreteType, + ctx: &Context, + ); + 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 +198,67 @@ 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 EnzymeTypeTreeInsertEq( + CTT: CTypeTreeRef, + indices: *const i64, + len: usize, + ct: CConcreteType, + ctx: &Context, + ) { + 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 +287,89 @@ 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 + } + + pub(crate) fn insert(&mut self, indices: &[i64], ct: CConcreteType, ctx: &Context) { + unsafe { + EnzymeTypeTreeInsertEq(self.inner, indices.as_ptr(), indices.len(), ct, ctx); + } + } +} + +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 { + <Self as std::fmt::Display>::fmt(self, f) + } +} + +impl Drop for TypeTree { + fn drop(&mut self) { + unsafe { EnzymeFreeTypeTree(self.inner) } + } +} diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs index 38a6a311954..afd2991a09c 100644 --- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs +++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs @@ -25,8 +25,8 @@ use rustc_target::spec::SymbolVisibility; use super::RustString; use super::debuginfo::{ DIArray, DIBuilder, DIDerivedType, DIDescriptor, DIEnumerator, DIFile, DIFlags, - DIGlobalVariableExpression, DILocation, DISPFlags, DIScope, DISubprogram, DISubrange, - DITemplateTypeParameter, DIType, DIVariable, DebugEmissionKind, DebugNameTableKind, + DIGlobalVariableExpression, DILocation, DISPFlags, DIScope, DISubprogram, + DITemplateTypeParameter, DIType, DebugEmissionKind, DebugNameTableKind, }; use crate::llvm; @@ -807,6 +807,8 @@ unsafe extern "C" { pub(crate) type Metadata; pub(crate) type BasicBlock; pub(crate) type Comdat; + /// `&'ll DbgRecord` represents `LLVMDbgRecordRef`. + pub(crate) type DbgRecord; } #[repr(C)] pub(crate) struct Builder<'a>(InvariantOpaque<'a>); @@ -891,7 +893,6 @@ pub(crate) mod debuginfo { pub(crate) type DIVariable = DIDescriptor; pub(crate) type DIGlobalVariableExpression = DIDescriptor; pub(crate) type DIArray = DIDescriptor; - pub(crate) type DISubrange = DIDescriptor; pub(crate) type DIEnumerator = DIDescriptor; pub(crate) type DITemplateTypeParameter = DIDescriptor; @@ -1992,6 +1993,59 @@ unsafe extern "C" { Scope: Option<&'ll Metadata>, AlignInBits: u32, // (optional; default is 0) ) -> &'ll Metadata; + + pub(crate) fn LLVMDIBuilderGetOrCreateSubrange<'ll>( + Builder: &DIBuilder<'ll>, + LowerBound: i64, + Count: i64, + ) -> &'ll Metadata; + + pub(crate) fn LLVMDIBuilderGetOrCreateArray<'ll>( + Builder: &DIBuilder<'ll>, + Data: *const Option<&'ll Metadata>, + NumElements: size_t, + ) -> &'ll Metadata; + + pub(crate) fn LLVMDIBuilderCreateExpression<'ll>( + Builder: &DIBuilder<'ll>, + Addr: *const u64, + Length: size_t, + ) -> &'ll Metadata; + + pub(crate) fn LLVMDIBuilderInsertDeclareRecordAtEnd<'ll>( + Builder: &DIBuilder<'ll>, + Storage: &'ll Value, + VarInfo: &'ll Metadata, + Expr: &'ll Metadata, + DebugLoc: &'ll Metadata, + Block: &'ll BasicBlock, + ) -> &'ll DbgRecord; + + pub(crate) fn LLVMDIBuilderCreateAutoVariable<'ll>( + Builder: &DIBuilder<'ll>, + Scope: &'ll Metadata, + Name: *const c_uchar, // See "PTR_LEN_STR". + NameLen: size_t, + File: &'ll Metadata, + LineNo: c_uint, + Ty: &'ll Metadata, + AlwaysPreserve: llvm::Bool, // "If true, this descriptor will survive optimizations." + Flags: DIFlags, + AlignInBits: u32, + ) -> &'ll Metadata; + + pub(crate) fn LLVMDIBuilderCreateParameterVariable<'ll>( + Builder: &DIBuilder<'ll>, + Scope: &'ll Metadata, + Name: *const c_uchar, // See "PTR_LEN_STR". + NameLen: size_t, + ArgNo: c_uint, + File: &'ll Metadata, + LineNo: c_uint, + Ty: &'ll Metadata, + AlwaysPreserve: llvm::Bool, // "If true, this descriptor will survive optimizations." + Flags: DIFlags, + ) -> &'ll Metadata; } #[link(name = "llvm-wrapper", kind = "static")] @@ -2358,43 +2412,6 @@ unsafe extern "C" { AlignInBits: u32, ) -> &'a DIGlobalVariableExpression; - pub(crate) fn LLVMRustDIBuilderCreateVariable<'a>( - Builder: &DIBuilder<'a>, - Tag: c_uint, - Scope: &'a DIDescriptor, - Name: *const c_char, - NameLen: size_t, - File: &'a DIFile, - LineNo: c_uint, - Ty: &'a DIType, - AlwaysPreserve: bool, - Flags: DIFlags, - ArgNo: c_uint, - AlignInBits: u32, - ) -> &'a DIVariable; - - pub(crate) fn LLVMRustDIBuilderGetOrCreateSubrange<'a>( - Builder: &DIBuilder<'a>, - Lo: i64, - Count: i64, - ) -> &'a DISubrange; - - pub(crate) fn LLVMRustDIBuilderGetOrCreateArray<'a>( - Builder: &DIBuilder<'a>, - Ptr: *const Option<&'a DIDescriptor>, - Count: c_uint, - ) -> &'a DIArray; - - pub(crate) fn LLVMRustDIBuilderInsertDeclareAtEnd<'a>( - Builder: &DIBuilder<'a>, - Val: &'a Value, - VarInfo: &'a DIVariable, - AddrOps: *const u64, - AddrOpsCount: c_uint, - DL: &'a DILocation, - InsertAtEnd: &'a BasicBlock, - ); - pub(crate) fn LLVMRustDIBuilderCreateEnumerator<'a>( Builder: &DIBuilder<'a>, Name: *const c_char, diff --git a/compiler/rustc_codegen_llvm/src/typetree.rs b/compiler/rustc_codegen_llvm/src/typetree.rs new file mode 100644 index 00000000000..7e263503700 --- /dev/null +++ b/compiler/rustc_codegen_llvm/src/typetree.rs @@ -0,0 +1,122 @@ +use rustc_ast::expand::typetree::FncTree; +#[cfg(feature = "llvm_enzyme")] +use { + crate::attributes, + rustc_ast::expand::typetree::TypeTree as RustTypeTree, + std::ffi::{CString, c_char, c_uint}, +}; + +use crate::llvm::{self, Value}; + +#[cfg(feature = "llvm_enzyme")] +fn to_enzyme_typetree( + rust_typetree: RustTypeTree, + _data_layout: &str, + llcx: &llvm::Context, +) -> llvm::TypeTree { + let mut enzyme_tt = llvm::TypeTree::new(); + process_typetree_recursive(&mut enzyme_tt, &rust_typetree, &[], llcx); + enzyme_tt +} +#[cfg(feature = "llvm_enzyme")] +fn process_typetree_recursive( + enzyme_tt: &mut llvm::TypeTree, + rust_typetree: &RustTypeTree, + parent_indices: &[i64], + llcx: &llvm::Context, +) { + 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, + }; + + let mut indices = parent_indices.to_vec(); + if !parent_indices.is_empty() { + indices.push(rust_type.offset as i64); + } else if rust_type.offset == -1 { + indices.push(-1); + } else { + indices.push(rust_type.offset as i64); + } + + enzyme_tt.insert(&indices, concrete_type, llcx); + + if rust_type.kind == rustc_ast::expand::typetree::Kind::Pointer + && !rust_type.child.0.is_empty() + { + process_typetree_recursive(enzyme_tt, &rust_type.child, &indices, llcx); + } + } +} + +#[cfg(feature = "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; + + 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"); + + let attr_name = "enzyme_type"; + let c_attr_name = CString::new(attr_name).unwrap(); + + for (i, input) in inputs.iter().enumerate() { + unsafe { + let enzyme_tt = to_enzyme_typetree(input.clone(), llvm_data_layout, llcx); + let c_str = llvm::EnzymeTypeTreeToString(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]); + llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr()); + } + } + + 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, + ); + + attributes::apply_to_llfn(fn_def, llvm::AttributePlace::ReturnValue, &[ret_attr]); + llvm::EnzymeTypeTreeToStringFree(c_str.as_ptr()); + } +} + +#[cfg(not(feature = "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 ab08125217f..d48c7cf874a 100644 --- a/compiler/rustc_codegen_llvm/src/va_arg.rs +++ b/compiler/rustc_codegen_llvm/src/va_arg.rs @@ -738,6 +738,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 { |