Simplify and dedup intrinsic call logic

Author: sayantn

compiler/rustc_codegen_llvm/src/allocator.rs

@@ -107,7 +107,7 @@ fn create_wrapper_function(
     let ty = cx.type_func(args, output.unwrap_or_else(|| cx.type_void()));
     let llfn = declare_simple_fn(
         &cx,
-        from_name,
+        from_name.as_bytes(),
         llvm::CallConv::CCallConv,
         llvm::UnnamedAddr::Global,
         llvm::Visibility::from_generic(tcx.sess.default_visibility()),
@@ -130,7 +130,7 @@ fn create_wrapper_function(
 
     let callee = declare_simple_fn(
         &cx,
-        to_name,
+        to_name.as_bytes(),
         llvm::CallConv::CCallConv,
         llvm::UnnamedAddr::Global,
         llvm::Visibility::Hidden,
@@ -150,7 +150,7 @@ fn create_wrapper_function(
         .enumerate()
         .map(|(i, _)| llvm::get_param(llfn, i as c_uint))
         .collect::<Vec<_>>();
-    let ret = bx.call(ty, callee, &args, None);
+    let ret = bx.simple_call(ty, callee, &args, None);
     llvm::LLVMSetTailCall(ret, True);
     if output.is_some() {
         bx.ret(ret);

compiler/rustc_codegen_llvm/src/builder.rs

@@ -31,8 +31,9 @@ use tracing::{debug, instrument, warn};
 
 use crate::abi::FnAbiLlvmExt;
 use crate::attributes;
-use crate::common::{AsCCharPtr, Funclet};
+use crate::common::Funclet;
 use crate::context::{CodegenCx, FullCx, GenericCx, SCx};
+use crate::declare::declare_simple_fn;
 use crate::llvm::{
     self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, GEPNoWrapFlags, Metadata, True,
 };
@@ -57,39 +58,6 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> Drop for GenericBuilder<'a, 'll, CX> {
     }
 }
 
-impl<'a, 'll> SBuilder<'a, 'll> {
-    pub(crate) fn call(
-        &mut self,
-        llty: &'ll Type,
-        llfn: &'ll Value,
-        args: &[&'ll Value],
-        funclet: Option<&Funclet<'ll>>,
-    ) -> &'ll Value {
-        debug!("call {:?} with args ({:?})", llfn, args);
-
-        let args = self.cast_arguments("call", llty, llfn, args, false);
-        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);
-        }
-
-        let call = unsafe {
-            llvm::LLVMBuildCallWithOperandBundles(
-                self.llbuilder,
-                llty,
-                llfn,
-                args.as_ptr(),
-                args.len() as c_uint,
-                bundles.as_ptr(),
-                bundles.len() as c_uint,
-                c"".as_ptr(),
-            )
-        };
-        call
-    }
-}
-
 impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
     fn with_cx(scx: &'a GenericCx<'ll, CX>) -> Self {
         // Create a fresh builder from the simple context.
@@ -1610,31 +1578,33 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
 
     pub(crate) fn simple_call(
         &mut self,
-        name: &[u8],
-        fn_ty: &'ll Type,
+        llty: &'ll Type,
+        llfn: &'ll Value,
         args: &[&'ll Value],
+        funclet: Option<&Funclet<'ll>>,
     ) -> &'ll Value {
-        let llfn = unsafe {
-            llvm::LLVMRustGetOrInsertFunction(
-                self.cx.llmod(),
-                name.as_c_char_ptr(),
-                name.len(),
-                fn_ty,
-            )
-        };
+        debug!("simple_call {:?} with args ({:?})", llfn, args);
 
-        let args = self.cast_arguments("simple_call", fn_ty, llfn, args, false);
+        let args = self.cast_arguments("simple_call", llty, llfn, args, false);
+        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);
+        }
 
-        unsafe {
-            llvm::LLVMBuildCall2(
+        let call = unsafe {
+            llvm::LLVMBuildCallWithOperandBundles(
                 self.llbuilder,
-                fn_ty,
+                llty,
                 llfn,
                 args.as_ptr(),
-                args.len() as _,
+                args.len() as c_uint,
+                bundles.as_ptr(),
+                bundles.len() as c_uint,
                 c"".as_ptr(),
             )
-        }
+        };
+        call
     }
 
     pub(crate) fn call_intrinsic(
@@ -1651,13 +1621,23 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
         let fn_ty = self.cx.intrinsic_type(intrinsic, type_params);
 
         let full_name = if intrinsic.is_overloaded() {
+            assert!(!type_params.is_empty());
             &intrinsic.overloaded_name(self.cx.llmod(), type_params)
         } else {
             assert!(type_params.is_empty());
             base_name.as_bytes()
         };
 
-        self.simple_call(full_name, fn_ty, args)
+        let llfn = declare_simple_fn(
+            self.cx,
+            full_name,
+            llvm::CCallConv,
+            llvm::UnnamedAddr::No,
+            llvm::Visibility::Default,
+            fn_ty,
+        );
+
+        self.simple_call(fn_ty, llfn, args, None)
     }
 }
 
@@ -1710,7 +1690,9 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         let fn_ty =
             self.type_func(&[self.type_ptr(), self.type_ptr(), self.type_isize()], llreturn_ty);
 
-        self.simple_call(b"memcmp", fn_ty, &[ptr1, ptr2, num])
+        let llfn = self.declare_cfn("memcmp", llvm::UnnamedAddr::No, fn_ty);
+
+        self.simple_call(fn_ty, llfn, &[ptr1, ptr2, num], None)
     }
 
     fn cast_return(

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -296,18 +296,17 @@ fn generate_enzyme_call<'ll>(
     attrs: AutoDiffAttrs,
 ) {
     // 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",
-        DiffMode::Reverse => "__enzyme_autodiff",
+    let mut ad_name = match attrs.mode {
+        DiffMode::Forward => b"__enzyme_fwddiff".as_slice(),
+        DiffMode::Reverse => b"__enzyme_autodiff",
         _ => panic!("logic bug in autodiff, unrecognized mode"),
     }
-    .to_string();
+    .to_vec();
 
     // add outer_fn name to ad_name to make it unique, in case users apply autodiff to multiple
-    // functions. Unwrap will only panic, if LLVM gave us an invalid string.
-    let name = llvm::get_value_name(outer_fn);
-    let outer_fn_name = std::str::from_utf8(name).unwrap();
-    ad_name.push_str(outer_fn_name);
+    // functions.
+    let outer_fn_name = llvm::get_value_name(outer_fn);
+    ad_name.extend_from_slice(outer_fn_name);
 
     // Let us assume the user wrote the following function square:
     //
@@ -400,7 +399,7 @@ fn generate_enzyme_call<'ll>(
             has_sret,
         );
 
-        let call = builder.call(enzyme_ty, ad_fn, &args, None);
+        let call = builder.simple_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 attached to it, but we just created this code oota. Given that the

compiler/rustc_codegen_llvm/src/declare.rs

@@ -24,7 +24,7 @@ use tracing::debug;
 
 use crate::abi::FnAbiLlvmExt;
 use crate::common::AsCCharPtr;
-use crate::context::{CodegenCx, GenericCx, SCx, SimpleCx};
+use crate::context::{CodegenCx, GenericCx, SCx};
 use crate::llvm::AttributePlace::Function;
 use crate::llvm::Visibility;
 use crate::type_::Type;
@@ -35,17 +35,17 @@ use crate::{attributes, llvm};
 ///
 /// If there’s a value with the same name already declared, the function will
 /// update the declaration and return existing Value instead.
-pub(crate) fn declare_simple_fn<'ll>(
-    cx: &SimpleCx<'ll>,
-    name: &str,
+pub(crate) fn declare_simple_fn<'ll, CX: Borrow<SCx<'ll>>>(
+    cx: &GenericCx<'ll, CX>,
+    name: &[u8],
     callconv: llvm::CallConv,
     unnamed: llvm::UnnamedAddr,
     visibility: llvm::Visibility,
     ty: &'ll Type,
 ) -> &'ll Value {
     debug!("declare_simple_fn(name={:?}, ty={:?})", name, ty);
     let llfn = unsafe {
-        llvm::LLVMRustGetOrInsertFunction(cx.llmod, name.as_c_char_ptr(), name.len(), ty)
+        llvm::LLVMRustGetOrInsertFunction(cx.llmod(), name.as_c_char_ptr(), name.len(), ty)
     };
 
     llvm::SetFunctionCallConv(llfn, callconv);
@@ -68,7 +68,7 @@ pub(crate) fn declare_raw_fn<'ll, 'tcx>(
     ty: &'ll Type,
 ) -> &'ll Value {
     debug!("declare_raw_fn(name={:?}, ty={:?})", name, ty);
-    let llfn = declare_simple_fn(cx, name, callconv, unnamed, visibility, ty);
+    let llfn = declare_simple_fn(cx, name.as_bytes(), callconv, unnamed, visibility, ty);
 
     let mut attrs = SmallVec::<[_; 4]>::new();
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1753,14 +1753,6 @@ unsafe extern "C" {
     ) -> *mut OperandBundle<'_>;
     pub(crate) fn LLVMDisposeOperandBundle(Bundle: ptr::NonNull<OperandBundle<'_>>);
 
-    pub(crate) fn LLVMBuildCall2<'a>(
-        B: &Builder<'a>,
-        Ty: &'a Type,
-        Fn: &'a Value,
-        Args: *const &'a Value,
-        NumArgs: c_uint,
-        Name: *const c_char,
-    ) -> &'a Value;
     pub(crate) fn LLVMBuildCallWithOperandBundles<'a>(
         B: &Builder<'a>,
         Ty: &'a Type,

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -356,7 +356,7 @@ impl Intrinsic {
         self,
         llmod: &'ll Module,
         type_params: &[&'ll Type],
-    ) -> CSlice<u8> {
+    ) -> Vec<u8> {
         let mut len = 0;
         let ptr = unsafe {
             LLVMIntrinsicCopyOverloadedName2(
@@ -367,39 +367,15 @@ impl Intrinsic {
                 &mut len,
             )
         };
-        unsafe { CSlice::from_raw(ptr.cast(), len) }
-    }
-}
-
-pub(crate) struct CSlice<T: Copy> {
-    ptr: *mut T,
-    len: usize,
-}
 
-impl<T: Copy> CSlice<T> {
-    pub(crate) unsafe fn from_raw(ptr: *mut T, len: usize) -> Self {
-        assert!(std::mem::size_of::<T>() != 0, "ZSTs aren't supported yet");
-        Self { ptr, len }
-    }
-}
+        let slice = unsafe { slice::from_raw_parts_mut(ptr.cast(), len) };
+        let copied = slice.to_vec();
 
-impl<T: Copy> std::ops::Deref for CSlice<T> {
-    type Target = [T];
-
-    fn deref(&self) -> &Self::Target {
-        unsafe { slice::from_raw_parts(self.ptr, self.len) }
-    }
-}
-
-impl<T: Copy> std::ops::DerefMut for CSlice<T> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        unsafe { slice::from_raw_parts_mut(self.ptr, self.len) }
-    }
-}
+        unsafe {
+            libc::free(ptr.cast());
+        }
 
-impl<T: Copy> Drop for CSlice<T> {
-    fn drop(&mut self) {
-        unsafe { libc::free(self.ptr.cast()) }
+        copied
     }
 }