wip

ZuseZ4 · ZuseZ4 · commit 55581c45fa54 · 2025-06-05T12:59:00.000-07:00
diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -653,6 +653,7 @@ pub(crate) fn run_pass_manager(
     // We then run the llvm_optimize function a second time, to optimize the code which we generated
     // in the enzyme differentiation pass.
     let enable_ad = config.autodiff.contains(&config::AutoDiff::Enable);
+    let enable_gpu = true;//config.offload.contains(&config::Offload::Enable);
     let stage = if thin {
         write::AutodiffStage::PreAD
     } else {
@@ -667,6 +668,114 @@ pub(crate) fn run_pass_manager(
         write::llvm_optimize(cgcx, dcx, module, None, config, opt_level, opt_stage, stage)?;
     }
 
+    if cfg!(llvm_enzyme) && enable_gpu && !thin {
+        // first we need to add all the fun to the host module
+        // %struct.__tgt_offload_entry = type { i64, i16, i16, i32, ptr, ptr, i64, i64, ptr }
+        // %struct.__tgt_kernel_arguments = type { i32, i32, ptr, ptr, ptr, ptr, ptr, ptr, i64, i64, [3 x i32], [3 x i32], i32 }
+        let cx =
+            SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
+        if cx.get_function("gen_tgt_offload").is_some() {
+            let offload_entry_ty = cx.type_named_struct("struct.__tgt_offload_entry");
+            let kernel_arguments_ty = cx.type_named_struct("struct.__tgt_kernel_arguments");
+            let tptr = cx.type_ptr();
+            let ti64 = cx.type_i64();
+            let ti32 = cx.type_i32();
+            let ti16 = cx.type_i16();
+            let tarr = cx.type_array(ti32, 3);
+
+            let entry_elements = vec![ti64, ti16, ti16, ti32, tptr, tptr, ti64, ti64, tptr];
+            let kernel_elements = vec![ti32, ti32, tptr, tptr, tptr, tptr, tptr, tptr, ti64, ti64, tarr, tarr, ti32];
+
+            cx.set_struct_body(offload_entry_ty, &entry_elements, false);
+            cx.set_struct_body(kernel_arguments_ty, &kernel_elements, false);
+            let global = cx.declare_global("my_struct_global", offload_entry_ty);
+            let global = cx.declare_global("my_struct_global2", kernel_arguments_ty);
+            dbg!(&offload_entry_ty);
+            dbg!(&kernel_arguments_ty);
+            //LLVMTypeRef elements[9] = {i64Ty, i16Ty, i16Ty, i32Ty, ptrTy, ptrTy, i64Ty, i64Ty, ptrTy};
+            //LLVMStructSetBody(structTy, elements, 9, 0);
+            dbg!("created struct");
+            for num in 0..5 {
+                if !cx.get_function(&format!("kernel_{num}")).is_some() {
+                    continue;
+                }
+            //for function in cx.get_functions() {
+                //if !attributes::has_attr(function, Function, llvm::AttributeKind::OptimizeForSize) {
+                //    dbg!("skipping minsize fnc");
+                //    dbg!(&function);
+                //    // print fnc name
+                //    let enzyme_marker = "minsize";
+                //    if attributes::has_string_attr(function, enzyme_marker) {
+                //        dbg!("found minsize str");
+                //    }
+                //    continue;
+
+                let size_name = format!(".offload_sizes.{num}");
+                let size_ty = cx.type_array(ti64, 4);
+                //let size_val = vec![8i64,0,16,0];
+                let c_val_8 = cx.get_const_i64(8);
+                let c_val_0 = cx.get_const_i64(0);
+                let c_val_16 = cx.get_const_i64(16);
+                let size_val = vec![c_val_8, c_val_0, c_val_16, c_val_0];
+
+                //let val = cx.define_global(&size_name, size_ty).unwrap();
+                //dbg!(&val);
+                //let section_var = cx
+                //    .define_global(section_var_name, llvm_type)
+                //    .unwrap_or_else(|| bug!("symbol `{}` is already defined", section_var_name));
+                //llvm::set_section(section_var, c".debug_gdb_scripts");
+                //llvm::set_initializer(section_var, cx.const_bytes(section_contents));
+                //llvm::LLVMSetGlobalConstant(section_var, llvm::True);
+                //llvm::set_linkage(section_var, llvm::Linkage::LinkOnceODRLinkage);
+                //// This should make sure that the whole section is not larger than
+                //// the string it contains. Otherwise we get a warning from GDB.
+                //llvm::LLVMSetAlignment(section_var, 1);
+                //llvm::set_initializer(val, cx.const_bytes(size_val.as_slice()));
+                let initializer = cx.const_array(ti64, &size_val);
+                let name = format!(".offload_sizes.{num}");
+                let c_name = CString::new(name).unwrap();
+                let array = llvm::add_global(cx.llmod, cx.val_ty(initializer), &c_name );
+                llvm::set_global_constant(array, true);
+                unsafe {llvm::LLVMSetUnnamedAddress(array, llvm::UnnamedAddr::Global)};
+                llvm::set_linkage(array, llvm::Linkage::PrivateLinkage);
+                llvm::set_initializer(array, initializer);
+                dbg!(&array);
+                // 1. @.offload_sizes.{num} = private unnamed_addr constant [4 x i64] [i64 8, i64 0, i64 16, i64 0]
+                // 2. @.offload_maptypes
+                // 3. @.__omp_offloading_<hash>_fnc_name_<hash> = weak constant i8 0
+                // 4. @.offloading.entry_name = internal unnamed_addr constant [66 x i8] c"__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7\00", section ".llvm.rodata.offloading", align 1
+                // 5. @.offloading.entry.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7.region_id, ptr @.offloading.entry_name, i64 0, i64 0, ptr null }, section "omp_offloading_entries", align 1
+            }
+            // @.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7.region_id = weak constant i8 0
+            // @.offload_sizes = private unnamed_addr constant [4 x i64] [i64 8, i64 0, i64 16, i64 0]
+            // @.offload_maptypes = private unnamed_addr constant [4 x i64] [i64 800, i64 544, i64 547, i64 544]
+            // @.__omp_offloading_86fafab6_c40006a1__Z3barPSt7complexIdES1_S0_m_l13.region_id = weak constant i8 0
+            // @.offload_sizes.1 = private unnamed_addr constant [4 x i64] [i64 8, i64 0, i64 16, i64 0]
+            // @.offload_maptypes.2 = private unnamed_addr constant [4 x i64] [i64 800, i64 544, i64 547, i64 544]
+            // @.__omp_offloading_86fafab6_c40006a1__Z5zaxpyPSt7complexIdES1_S0_m_l19.region_id = weak constant i8 0
+            // @.offload_sizes.3 = private unnamed_addr constant [4 x i64] [i64 8, i64 0, i64 16, i64 0]
+            // @.offload_maptypes.4 = private unnamed_addr constant [4 x i64] [i64 800, i64 544, i64 547, i64 544]
+            // @.offload_sizes.5 = private unnamed_addr constant [2 x i64] [i64 16384, i64 16384]
+            // @.offload_maptypes.6 = private unnamed_addr constant [2 x i64] [i64 1, i64 3]
+            // @_ZSt4cout = external global %"class.std::basic_ostream", align 8
+            // @.str = private unnamed_addr constant [3 x i8] c"hi\00", align 1
+            // @.offload_sizes.7 = private unnamed_addr constant [2 x i64] [i64 16384, i64 16384]
+            // @.offload_maptypes.8 = private unnamed_addr constant [2 x i64] [i64 1, i64 3]
+            // @.str.9 = private unnamed_addr constant [3 x i8] c"ho\00", align 1
+            // @.offloading.entry_name = internal unnamed_addr constant [66 x i8] c"__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7\00", section ".llvm.rodata.offloading", align 1
+            // @.offloading.entry.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.__omp_offloading_86fafab6_c40006a1__Z3fooPSt7complexIdES1_S0_m_l7.region_id, ptr @.offloading.entry_name, i64 0, i64 0, ptr null }, section "omp_offloading_entries", align 1
+            // @.offloading.entry_name.10 = internal unnamed_addr constant [67 x i8] c"__omp_offloading_86fafab6_c40006a1__Z3barPSt7complexIdES1_S0_m_l13\00", section ".llvm.rodata.offloading", align 1
+            // @.offloading.entry.__omp_offloading_86fafab6_c40006a1__Z3barPSt7complexIdES1_S0_m_l13 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.__omp_offloading_86fafab6_c40006a1__Z3barPSt7complexIdES1_S0_m_l13.region_id, ptr @.offloading.entry_name.10, i64 0, i64 0, ptr null }, section "omp_offloading_entries", align 1
+            // @.offloading.entry_name.11 = internal unnamed_addr constant [69 x i8] c"__omp_offloading_86fafab6_c40006a1__Z5zaxpyPSt7complexIdES1_S0_m_l19\00", section ".llvm.rodata.offloading", align 1
+            // @.offloading.entry.__omp_offloading_86fafab6_c40006a1__Z5zaxpyPSt7complexIdES1_S0_m_l19 = weak constant %struct.__tgt_offload_entry { i64 0, i16 1, i16 1, i32 0, ptr @.__omp_offloading_86fafab6_c40006a1__Z5zaxpyPSt7complexIdES1_S0_m_l19.region_id, ptr @.offloading.entry_name.11, i64 0, i64 0, ptr null }, section "omp_offloading_entries", align 1
+            // @llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @_GLOBAL__sub_I_zaxpy.cpp, ptr null }]
+        } else {
+            dbg!("no marker found");
+        }
+    } else {
+        dbg!("Not creating struct");
+    }
+
     if cfg!(llvm_enzyme) && enable_ad && !thin {
         let cx =
             SimpleCx::new(module.module_llvm.llmod(), &module.module_llvm.llcx, cgcx.pointer_size);
diff --git a/compiler/rustc_codegen_llvm/src/common.rs b/compiler/rustc_codegen_llvm/src/common.rs
@@ -99,14 +99,14 @@ impl<'ll, CX: Borrow<SCx<'ll>>> BackendTypes for GenericCx<'ll, CX> {
     type DIVariable = &'ll llvm::debuginfo::DIVariable;
 }
 
-impl<'ll> CodegenCx<'ll, '_> {
+impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
     pub(crate) fn const_array(&self, ty: &'ll Type, elts: &[&'ll Value]) -> &'ll Value {
         let len = u64::try_from(elts.len()).expect("LLVMConstArray2 elements len overflow");
         unsafe { llvm::LLVMConstArray2(ty, elts.as_ptr(), len) }
     }
 
     pub(crate) fn const_bytes(&self, bytes: &[u8]) -> &'ll Value {
-        bytes_in_context(self.llcx, bytes)
+        bytes_in_context(self.llcx(), bytes)
     }
 
     pub(crate) fn const_get_elt(&self, v: &'ll Value, idx: u64) -> &'ll Value {
diff --git a/compiler/rustc_codegen_llvm/src/declare.rs b/compiler/rustc_codegen_llvm/src/declare.rs
@@ -99,6 +99,43 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             )
         }
     }
+
+    /// Gets declared value by name.
+    pub(crate) fn get_declared_value(&self, name: &str) -> Option<&'ll Value> {
+        debug!("get_declared_value(name={:?})", name);
+        unsafe { llvm::LLVMRustGetNamedValue(self.llmod(), name.as_c_char_ptr(), name.len()) }
+    }
+
+    /// Gets defined or externally defined (AvailableExternally linkage) value by
+    /// name.
+    pub(crate) fn get_defined_value(&self, name: &str) -> Option<&'ll Value> {
+        self.get_declared_value(name).and_then(|val| {
+            let declaration = llvm::is_declaration(val);
+            if !declaration { Some(val) } else { None }
+        })
+    }
+
+    /// Declare a global with an intention to define it.
+    ///
+    /// Use this function when you intend to define a global. This function will
+    /// return `None` if the name already has a definition associated with it. In that
+    /// case an error should be reported to the user, because it usually happens due
+    /// to user’s fault (e.g., misuse of `#[no_mangle]` or `#[export_name]` attributes).
+    pub(crate) fn define_global(&self, name: &str, ty: &'ll Type) -> Option<&'ll Value> {
+        if self.get_defined_value(name).is_some() {
+            None
+        } else {
+            Some(self.declare_global(name, ty))
+        }
+    }
+
+    /// Declare a private global
+    ///
+    /// Use this function when you intend to define a global without a name.
+    pub(crate) fn define_private_global(&self, ty: &'ll Type) -> &'ll Value {
+        unsafe { llvm::LLVMRustInsertPrivateGlobal(self.llmod(), ty) }
+    }
+
 }
 
 impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
@@ -215,40 +252,4 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 
         llfn
     }
-
-    /// Declare a global with an intention to define it.
-    ///
-    /// Use this function when you intend to define a global. This function will
-    /// return `None` if the name already has a definition associated with it. In that
-    /// case an error should be reported to the user, because it usually happens due
-    /// to user’s fault (e.g., misuse of `#[no_mangle]` or `#[export_name]` attributes).
-    pub(crate) fn define_global(&self, name: &str, ty: &'ll Type) -> Option<&'ll Value> {
-        if self.get_defined_value(name).is_some() {
-            None
-        } else {
-            Some(self.declare_global(name, ty))
-        }
-    }
-
-    /// Declare a private global
-    ///
-    /// Use this function when you intend to define a global without a name.
-    pub(crate) fn define_private_global(&self, ty: &'ll Type) -> &'ll Value {
-        unsafe { llvm::LLVMRustInsertPrivateGlobal(self.llmod, ty) }
-    }
-
-    /// Gets declared value by name.
-    pub(crate) fn get_declared_value(&self, name: &str) -> Option<&'ll Value> {
-        debug!("get_declared_value(name={:?})", name);
-        unsafe { llvm::LLVMRustGetNamedValue(self.llmod, name.as_c_char_ptr(), name.len()) }
-    }
-
-    /// Gets defined or externally defined (AvailableExternally linkage) value by
-    /// name.
-    pub(crate) fn get_defined_value(&self, name: &str) -> Option<&'ll Value> {
-        self.get_declared_value(name).and_then(|val| {
-            let declaration = llvm::is_declaration(val);
-            if !declaration { Some(val) } else { None }
-        })
-    }
 }
diff --git a/compiler/rustc_codegen_ssa/src/back/write.rs b/compiler/rustc_codegen_ssa/src/back/write.rs
@@ -121,6 +121,7 @@ pub struct ModuleConfig {
     pub emit_lifetime_markers: bool,
     pub llvm_plugins: Vec<String>,
     pub autodiff: Vec<config::AutoDiff>,
+    pub offload: Vec<config::Offload>,
 }
 
 impl ModuleConfig {
@@ -270,6 +271,7 @@ impl ModuleConfig {
             emit_lifetime_markers: sess.emit_lifetime_markers(),
             llvm_plugins: if_regular!(sess.opts.unstable_opts.llvm_plugins.clone(), vec![]),
             autodiff: if_regular!(sess.opts.unstable_opts.autodiff.clone(), vec![]),
+            offload: if_regular!(sess.opts.unstable_opts.offload.clone(), vec![]),
         }
     }
 
diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs
@@ -226,6 +226,13 @@ pub enum CoverageLevel {
     Mcdc,
 }
 
+// The different settings that the `-Z offload` flag can have.
+#[derive(Clone, Copy, PartialEq, Hash, Debug)]
+pub enum Offload {
+    /// Enable the llvm offload pipeline
+    Enable,
+}
+
 /// The different settings that the `-Z autodiff` flag can have.
 #[derive(Clone, Copy, PartialEq, Hash, Debug)]
 pub enum AutoDiff {
@@ -3061,7 +3068,7 @@ pub(crate) mod dep_tracking {
     };
 
     use super::{
-        AutoDiff, BranchProtection, CFGuard, CFProtection, CollapseMacroDebuginfo, CoverageOptions,
+        AutoDiff, Offload, BranchProtection, CFGuard, CFProtection, CollapseMacroDebuginfo, CoverageOptions,
         CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug, FunctionReturn,
         InliningThreshold, InstrumentCoverage, InstrumentXRay, LinkerPluginLto, LocationDetail,
         LtoCli, MirStripDebugInfo, NextSolverConfig, OomStrategy, OptLevel, OutFileName,
@@ -3110,6 +3117,7 @@ pub(crate) mod dep_tracking {
 
     impl_dep_tracking_hash_via_hash!(
         AutoDiff,
+        Offload,
         bool,
         usize,
         NonZero<usize>,
diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs
@@ -712,6 +712,7 @@ mod desc {
     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`, `PrintAA`, `PrintPerf`, `PrintModBefore`, `PrintModAfter`, `PrintModFinal`, `PrintPasses`, `NoPostopt`, `LooseTypes`, `Inline`";
+    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;
     pub(crate) const parse_number: &str = "a number";
@@ -1343,6 +1344,27 @@ pub mod parse {
         }
     }
 
+    pub(crate) fn parse_offload(slot: &mut Vec<Offload>, v: Option<&str>) -> bool {
+        let Some(v) = v else {
+            *slot = vec![];
+            return true;
+        };
+        let mut v: Vec<&str> = v.split(",").collect();
+        v.sort_unstable();
+        for &val in v.iter() {
+            let variant = match val {
+                "Enable" => Offload::Enable,
+                _ => {
+                    // FIXME(ZuseZ4): print an error saying which value is not recognized
+                    return false;
+                }
+            };
+            slot.push(variant);
+        }
+
+        true
+    }
+
     pub(crate) fn parse_autodiff(slot: &mut Vec<AutoDiff>, v: Option<&str>) -> bool {
         let Some(v) = v else {
             *slot = vec![];
@@ -2372,6 +2394,11 @@ options! {
         "do not use unique names for text and data sections when -Z function-sections is used"),
     normalize_docs: bool = (false, parse_bool, [TRACKED],
         "normalize associated items in rustdoc when generating documentation"),
+    offload: Vec<crate::config::Offload> = (Vec::new(), parse_offload, [TRACKED],
+        "a list of offload flags to enable
+        Mandatory setting:
+        `=Enable`
+        Currently the only option available"),
     on_broken_pipe: OnBrokenPipe = (OnBrokenPipe::Default, parse_on_broken_pipe, [TRACKED],
         "behavior of std::io::ErrorKind::BrokenPipe (SIGPIPE)"),
     oom: OomStrategy = (OomStrategy::Abort, parse_oom_strategy, [TRACKED],

Original file line number	Diff line number	Diff line change
`@@ -99,14 +99,14 @@ impl<'ll, CX: Borrow<SCx<'ll>>> BackendTypes for GenericCx<'ll, CX> {`
`99`	`99`	`type DIVariable = &'ll llvm::debuginfo::DIVariable;`
`100`	`100`	`}`
`101`	`101`
`102`		`-impl<'ll> CodegenCx<'ll, '_> {`
	`102`	`+impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {`
`103`	`103`	`pub(crate) fn const_array(&self, ty: &'ll Type, elts: &[&'ll Value]) -> &'ll Value {`
`104`	`104`	`let len = u64::try_from(elts.len()).expect("LLVMConstArray2 elements len overflow");`
`105`	`105`	`unsafe { llvm::LLVMConstArray2(ty, elts.as_ptr(), len) }`
`106`	`106`	`}`
`107`	`107`
`108`	`108`	`pub(crate) fn const_bytes(&self, bytes: &[u8]) -> &'ll Value {`
`109`		`- bytes_in_context(self.llcx, bytes)`
	`109`	`+ bytes_in_context(self.llcx(), bytes)`
`110`	`110`	`}`
`111`	`111`
`112`	`112`	`pub(crate) fn const_get_elt(&self, v: &'ll Value, idx: u64) -> &'ll Value {`
Original file line number	Diff line number	Diff line change
`@@ -121,6 +121,7 @@ pub struct ModuleConfig {`
`121`	`121`	`pub emit_lifetime_markers: bool,`
`122`	`122`	`pub llvm_plugins: Vec<String>,`
`123`	`123`	`pub autodiff: Vec<config::AutoDiff>,`
	`124`	`+ pub offload: Vec<config::Offload>,`
`124`	`125`	`}`
`125`	`126`
`126`	`127`	`impl ModuleConfig {`
`@@ -270,6 +271,7 @@ impl ModuleConfig {`
`270`	`271`	`emit_lifetime_markers: sess.emit_lifetime_markers(),`
`271`	`272`	`llvm_plugins: if_regular!(sess.opts.unstable_opts.llvm_plugins.clone(), vec![]),`
`272`	`273`	`autodiff: if_regular!(sess.opts.unstable_opts.autodiff.clone(), vec![]),`
	`274`	`+ offload: if_regular!(sess.opts.unstable_opts.offload.clone(), vec![]),`
`273`	`275`	`}`
`274`	`276`	`}`
`275`	`277`