wip

Author: ZuseZ4

Diff for: compiler/rustc_codegen_llvm/src/builder.rs

@@ -31,12 +31,151 @@ use tracing::{debug, instrument};
 use crate::abi::FnAbiLlvmExt;
 use crate::attributes;
 use crate::common::Funclet;
-use crate::context::CodegenCx;
+use crate::context::{CodegenCx, SimpleCx};
 use crate::llvm::{self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, True};
 use crate::type_::Type;
 use crate::type_of::LayoutLlvmExt;
 use crate::value::Value;
 
+// All Builders must have an llfn associated with them
+#[must_use]
+pub(crate) struct SBuilder<'a, 'll> {
+    pub llbuilder: &'ll mut llvm::Builder<'ll>,
+    pub cx: &'a SimpleCx<'ll>,
+}
+
+impl Drop for SBuilder<'_, '_> {
+    fn drop(&mut self) {
+        unsafe {
+            llvm::LLVMDisposeBuilder(&mut *(self.llbuilder as *mut _));
+        }
+    }
+}
+
+impl<'a, 'll> SBuilder<'a, 'll> {
+    fn build(cx: &'a SimpleCx<'ll>, llbb: &'ll BasicBlock) -> SBuilder<'a, 'll> {
+        let bx = SBuilder::with_scx(cx);
+        unsafe {
+            llvm::LLVMPositionBuilderAtEnd(bx.llbuilder, llbb);
+        }
+        bx
+    }
+}
+
+impl<'a, 'll> SBuilder<'a, 'll> {
+    fn with_scx(scx: &'a SimpleCx<'ll>) -> Self {
+        // Create a fresh builder from the simple context.
+        let llbuilder = unsafe { llvm::LLVMCreateBuilderInContext(scx.llcx) };
+        SBuilder { llbuilder, cx: scx }
+    }
+
+    pub(crate) fn bitcast(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
+        unsafe { llvm::LLVMBuildBitCast(self.llbuilder, val, dest_ty, UNNAMED) }
+    }
+
+    fn ret_void(&mut self) {
+        unsafe {
+            llvm::LLVMBuildRetVoid(self.llbuilder);
+        }
+    }
+
+    fn ret(&mut self, v: &'ll Value) {
+        unsafe {
+            llvm::LLVMBuildRet(self.llbuilder, v);
+        }
+    }
+
+    fn check_call<'b>(
+        &mut self,
+        typ: &str,
+        fn_ty: &'ll Type,
+        llfn: &'ll Value,
+        args: &'b [&'ll Value],
+    ) -> Cow<'b, [&'ll Value]> {
+        assert!(
+            self.cx.type_kind(fn_ty) == TypeKind::Function,
+            "builder::{typ} not passed a function, but {fn_ty:?}"
+        );
+
+        let param_tys = self.cx.func_params_types(fn_ty);
+
+        let all_args_match = iter::zip(&param_tys, args.iter().map(|&v| self.cx.val_ty(v)))
+            .all(|(expected_ty, actual_ty)| *expected_ty == actual_ty);
+
+        if all_args_match {
+            return Cow::Borrowed(args);
+        }
+
+        let casted_args: Vec<_> = iter::zip(param_tys, args)
+            .enumerate()
+            .map(|(i, (expected_ty, &actual_val))| {
+                let actual_ty = self.cx.val_ty(actual_val);
+                if expected_ty != actual_ty {
+                    debug!(
+                        "type mismatch in function call of {:?}. \
+                            Expected {:?} for param {}, got {:?}; injecting bitcast",
+                        llfn, expected_ty, i, actual_ty
+                    );
+                    self.bitcast(actual_val, expected_ty)
+                } else {
+                    actual_val
+                }
+            })
+            .collect();
+
+        Cow::Owned(casted_args)
+    }
+
+    fn call(
+        &mut self,
+        llty: &'ll Type,
+        //fn_attrs: Option<&CodegenFnAttrs>,
+        //fn_abi: Option<&FnAbi<'tcx, Ty<'tcx>>>,
+        llfn: &'ll Value,
+        args: &[&'ll Value],
+        funclet: Option<&Funclet<'ll>>,
+        //instance: Option<Instance<'tcx>>,
+    ) -> &'ll Value {
+        debug!("call {:?} with args ({:?})", llfn, args);
+
+        let args = self.check_call("call", llty, llfn, args);
+        let funclet_bundle = funclet.map(|funclet| funclet.bundle());
+        let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
+        if let Some(funclet_bundle) = funclet_bundle {
+            bundles.push(funclet_bundle);
+        }
+
+        // Emit CFI pointer type membership test
+        //self.cfi_type_test(fn_attrs, fn_abi, instance, llfn);
+
+        // Emit KCFI operand bundle
+        //let kcfi_bundle = self.kcfi_operand_bundle(fn_attrs, fn_abi, instance, llfn);
+        //if let Some(kcfi_bundle) = kcfi_bundle.as_deref() {
+        //    bundles.push(kcfi_bundle);
+        //}
+
+        let call = unsafe {
+            llvm::LLVMBuildCallWithOperandBundles(
+                self.llbuilder,
+                llty,
+                llfn,
+                args.as_ptr() as *const &llvm::Value,
+                args.len() as c_uint,
+                bundles.as_ptr(),
+                bundles.len() as c_uint,
+                c"".as_ptr(),
+            )
+        };
+        //if let Some(fn_abi) = fn_abi {
+        //    fn_abi.apply_attrs_callsite(self, call);
+        //}
+        call
+    }
+
+}
+
+
+
 // All Builders must have an llfn associated with them
 #[must_use]
 pub(crate) struct Builder<'a, 'll, 'tcx> {
@@ -55,7 +194,7 @@ impl Drop for Builder<'_, '_, '_> {
 /// Empty string, to be used where LLVM expects an instruction name, indicating
 /// that the instruction is to be left unnamed (i.e. numbered, in textual IR).
 // FIXME(eddyb) pass `&CStr` directly to FFI once it's a thin pointer.
-const UNNAMED: *const c_char = c"".as_ptr();
+pub(crate) const UNNAMED: *const c_char = c"".as_ptr();
 
 impl<'ll, 'tcx> BackendTypes for Builder<'_, 'll, 'tcx> {
     type Value = <CodegenCx<'ll, 'tcx> as BackendTypes>::Value;
@@ -170,7 +309,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     fn append_block(cx: &'a CodegenCx<'ll, 'tcx>, llfn: &'ll Value, name: &str) -> &'ll BasicBlock {
         unsafe {
             let name = SmallCStr::new(name);
-            llvm::LLVMAppendBasicBlockInContext(cx.llcx, llfn, name.as_ptr())
+            llvm::LLVMAppendBasicBlockInContext(cx.scx.llcx, llfn, name.as_ptr())
         }
     }
 
@@ -621,14 +760,14 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                 llvm::LLVMValueAsMetadata(self.cx.const_uint_big(llty, range.start)),
                 llvm::LLVMValueAsMetadata(self.cx.const_uint_big(llty, range.end.wrapping_add(1))),
             ];
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, md.as_ptr(), md.len());
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, md.as_ptr(), md.len());
             self.set_metadata(load, llvm::MD_range, md);
         }
     }
 
     fn nonnull_metadata(&mut self, load: &'ll Value) {
         unsafe {
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, ptr::null(), 0);
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, ptr::null(), 0);
             self.set_metadata(load, llvm::MD_nonnull, md);
         }
     }
@@ -678,7 +817,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
                     //
                     // [1]: https://llvm.org/docs/LangRef.html#store-instruction
                     let one = llvm::LLVMValueAsMetadata(self.cx.const_i32(1));
-                    let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, &one, 1);
+                    let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, &one, 1);
                     self.set_metadata(store, llvm::MD_nontemporal, md);
                 }
             }
@@ -1144,7 +1283,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn set_invariant_load(&mut self, load: &'ll Value) {
         unsafe {
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, ptr::null(), 0);
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, ptr::null(), 0);
             self.set_metadata(load, llvm::MD_invariant_load, md);
         }
     }
@@ -1209,7 +1348,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn apply_attrs_to_cleanup_callsite(&mut self, llret: &'ll Value) {
         // Cleanup is always the cold path.
-        let cold_inline = llvm::AttributeKind::Cold.create_attr(self.llcx);
+        let cold_inline = llvm::AttributeKind::Cold.create_attr(self.scx.llcx);
         attributes::apply_to_callsite(llret, llvm::AttributePlace::Function, &[cold_inline]);
     }
 }
@@ -1224,7 +1363,7 @@ impl<'ll> StaticBuilderMethods for Builder<'_, 'll, '_> {
 impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
     fn with_cx(cx: &'a CodegenCx<'ll, 'tcx>) -> Self {
         // Create a fresh builder from the crate context.
-        let llbuilder = unsafe { llvm::LLVMCreateBuilderInContext(cx.llcx) };
+        let llbuilder = unsafe { llvm::LLVMCreateBuilderInContext(cx.scx.llcx) };
         Builder { llbuilder, cx }
     }
 
@@ -1241,21 +1380,21 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
     fn align_metadata(&mut self, load: &'ll Value, align: Align) {
         unsafe {
             let md = [llvm::LLVMValueAsMetadata(self.cx.const_u64(align.bytes()))];
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, md.as_ptr(), md.len());
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, md.as_ptr(), md.len());
             self.set_metadata(load, llvm::MD_align, md);
         }
     }
 
     fn noundef_metadata(&mut self, load: &'ll Value) {
         unsafe {
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, ptr::null(), 0);
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, ptr::null(), 0);
             self.set_metadata(load, llvm::MD_noundef, md);
         }
     }
 
     pub(crate) fn set_unpredictable(&mut self, inst: &'ll Value) {
         unsafe {
-            let md = llvm::LLVMMDNodeInContext2(self.cx.llcx, ptr::null(), 0);
+            let md = llvm::LLVMMDNodeInContext2(self.cx.scx.llcx, ptr::null(), 0);
             self.set_metadata(inst, llvm::MD_unpredictable, md);
         }
     }

Diff for: compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -3,20 +3,22 @@ use std::ptr;
 use rustc_ast::expand::autodiff_attrs::{AutoDiffAttrs, AutoDiffItem, DiffActivity, DiffMode};
 use rustc_codegen_ssa::ModuleCodegen;
 use rustc_codegen_ssa::back::write::ModuleConfig;
-use rustc_codegen_ssa::traits::{BaseTypeCodegenMethods, BuilderMethods};
+//use rustc_codegen_ssa::traits::BuilderMethods;
 use rustc_errors::FatalError;
-use rustc_middle::ty::TyCtxt;
+//use rustc_middle::ty::TyCtxt;
 use rustc_session::config::Lto;
 use tracing::{debug, trace};
 
 use crate::back::write::{llvm_err, llvm_optimize};
-use crate::builder::Builder;
-use crate::declare::declare_raw_fn;
+use crate::builder::SBuilder;
+//use crate::builder::{Builder, SBuilder};
+use crate::context::SimpleCx;
+use crate::declare::declare_simple_fn;
 use crate::errors::LlvmError;
 use crate::llvm::AttributePlace::Function;
 use crate::llvm::{Metadata, True};
 use crate::value::Value;
-use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, context, llvm};
+use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
 
 fn get_params(fnc: &Value) -> Vec<&Value> {
     unsafe {
@@ -38,8 +40,8 @@ fn get_params(fnc: &Value) -> Vec<&Value> {
 /// [^1]: <https://enzyme.mit.edu/getting_started/CallingConvention/>
 // FIXME(ZuseZ4): `outer_fn` should include upstream safety checks to
 // cover some assumptions of enzyme/autodiff, which could lead to UB otherwise.
-fn generate_enzyme_call<'ll, 'tcx>(
-    cx: &context::CodegenCx<'ll, 'tcx>,
+fn generate_enzyme_call<'ll>(
+    cx: &SimpleCx<'ll>,
     fn_to_diff: &'ll Value,
     outer_fn: &'ll Value,
     attrs: AutoDiffAttrs,
@@ -112,7 +114,7 @@ fn generate_enzyme_call<'ll, 'tcx>(
         //FIXME(ZuseZ4): the CC/Addr/Vis values are best effort guesses, we should look at tests and
         // think a bit more about what should go here.
         let cc = llvm::LLVMGetFunctionCallConv(outer_fn);
-        let ad_fn = declare_raw_fn(
+        let ad_fn = declare_simple_fn(
             cx,
             &ad_name,
             llvm::CallConv::try_from(cc).expect("invalid callconv"),
@@ -132,7 +134,7 @@ fn generate_enzyme_call<'ll, 'tcx>(
         llvm::LLVMRustEraseInstFromParent(br);
 
         let last_inst = llvm::LLVMRustGetLastInstruction(entry).unwrap();
-        let mut builder = Builder::build(cx, entry);
+        let mut builder = SBuilder::build(cx, entry);
 
         let num_args = llvm::LLVMCountParams(&fn_to_diff);
         let mut args = Vec::with_capacity(num_args as usize + 1);
@@ -236,7 +238,7 @@ fn generate_enzyme_call<'ll, 'tcx>(
             }
         }
 
-        let call = builder.call(enzyme_ty, None, None, ad_fn, &args, None, None);
+        let call = builder.call(enzyme_ty, ad_fn, &args, None);
 
         // This part is a bit iffy. LLVM requires that a call to an inlineable function has some
         // metadata attachted to it, but we just created this code oota. Given that the
@@ -274,20 +276,23 @@ fn generate_enzyme_call<'ll, 'tcx>(
     }
 }
 
-pub(crate) fn differentiate<'ll, 'tcx>(
+pub(crate) fn differentiate<'ll>(
     module: &'ll ModuleCodegen<ModuleLlvm>,
     cgcx: &CodegenContext<LlvmCodegenBackend>,
-    tcx: TyCtxt<'tcx>,
+    //cx: SimpleCx<'ll>,
     diff_items: Vec<AutoDiffItem>,
     config: &ModuleConfig,
 ) -> Result<(), FatalError> {
     for item in &diff_items {
         trace!("{}", item);
     }
 
+    let cx = SimpleCx { llmod: module.module_llvm.llmod(), llcx: module.module_llvm.llcx };
+
     let diag_handler = cgcx.create_dcx();
-    let (_, cgus) = tcx.collect_and_partition_mono_items(());
-    let cx = context::CodegenCx::new(tcx, &cgus.first().unwrap(), &module.module_llvm);
+    //let cx = context::SimpleCx { llmod: module.module_llvm.llmod(), llcx: module.module_llvm.llcx };
+    //let (_, cgus) = tcx.collect_and_partition_mono_items(());
+    //let cx = context::CodegenCx::new(tcx, &cgus.first().unwrap(), &module.module_llvm);
 
     // Before dumping the module, we want all the TypeTrees to become part of the module.
     for item in diff_items.iter() {