Rollup merge of #148533 - bjorn3:split_llvm_intrinsic_abi_handling, r=WaffleLapkin

Split LLVM intrinsic abi handling from the rest of the abi handling

LLVM intrinsics have weird requirements like requiring the fake "unadjusted" abi, not being callable through function pointers and for all codegen backends other than cg_llvm requiring special cases to redirect them to the correct backend specific intrinsic (or directly codegen their implementation inline without any intrinsic call). By splitting the LLVM intrinsic handling it becomes easier for backends to special case them and should in the future allow getting rid of the abi calculation for `extern "unadjusted"` in favor of computing the correct abi directly in the backend without depending on the exact way cg_ssa lowers types.

compiler/rustc_codegen_gcc/src/builder.rs

@@ -314,14 +314,12 @@ pub fn current_func(&self) -> Function<'gcc> {
         self.block.get_function()
     }
 
-    fn function_call(
+    pub fn function_call(
         &mut self,
-        func: RValue<'gcc>,
+        func: Function<'gcc>,
         args: &[RValue<'gcc>],
         _funclet: Option<&Funclet>,
     ) -> RValue<'gcc> {
-        // TODO(antoyo): remove when the API supports a different type for functions.
-        let func: Function<'gcc> = self.cx.rvalue_as_function(func);
         let args = self.check_call("call", func, args);
 
         // gccjit requires to use the result of functions, even when it's not used.
@@ -514,6 +512,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
     type CodegenCx = CodegenCx<'gcc, 'tcx>;
 
     fn build(cx: &'a CodegenCx<'gcc, 'tcx>, block: Block<'gcc>) -> Builder<'a, 'gcc, 'tcx> {
+        *cx.current_func.borrow_mut() = Some(block.get_function());
         Builder::with_cx(cx, block)
     }
 
@@ -1765,6 +1764,8 @@ fn call(
         // FIXME(antoyo): remove when having a proper API.
         let gcc_func = unsafe { std::mem::transmute::<RValue<'gcc>, Function<'gcc>>(func) };
         let call = if self.functions.borrow().values().any(|value| *value == gcc_func) {
+            // TODO(antoyo): remove when the API supports a different type for functions.
+            let func: Function<'gcc> = self.cx.rvalue_as_function(func);
             self.function_call(func, args, funclet)
         } else {
             // If it's a not function that was defined, it's a function pointer.

compiler/rustc_codegen_gcc/src/context.rs

@@ -92,6 +92,8 @@ pub struct CodegenCx<'gcc, 'tcx> {
     pub instances: RefCell<FxHashMap<Instance<'tcx>, LValue<'gcc>>>,
     /// Cache function instances of monomorphic and polymorphic items
     pub function_instances: RefCell<FxHashMap<Instance<'tcx>, Function<'gcc>>>,
+    /// Cache function instances of intrinsics
+    pub intrinsic_instances: RefCell<FxHashMap<Instance<'tcx>, Function<'gcc>>>,
     /// Cache generated vtables
     pub vtables:
         RefCell<FxHashMap<(Ty<'tcx>, Option<ty::ExistentialTraitRef<'tcx>>), RValue<'gcc>>>,
@@ -280,6 +282,7 @@ pub fn new(
             linkage: Cell::new(FunctionType::Internal),
             instances: Default::default(),
             function_instances: Default::default(),
+            intrinsic_instances: Default::default(),
             on_stack_params: Default::default(),
             on_stack_function_params: Default::default(),
             vtables: Default::default(),
@@ -391,17 +394,13 @@ fn vtables(
     }
 
     fn get_fn(&self, instance: Instance<'tcx>) -> Function<'gcc> {
-        let func = get_fn(self, instance);
-        *self.current_func.borrow_mut() = Some(func);
-        func
+        get_fn(self, instance)
     }
 
     fn get_fn_addr(&self, instance: Instance<'tcx>) -> RValue<'gcc> {
         let func_name = self.tcx.symbol_name(instance).name;
 
-        let func = if self.intrinsics.borrow().contains_key(func_name) {
-            self.intrinsics.borrow()[func_name]
-        } else if let Some(variable) = self.get_declared_value(func_name) {
+        let func = if let Some(variable) = self.get_declared_value(func_name) {
             return variable;
         } else {
             get_fn(self, instance)

compiler/rustc_codegen_gcc/src/declare.rs

@@ -8,7 +8,6 @@
 
 use crate::abi::{FnAbiGcc, FnAbiGccExt};
 use crate::context::CodegenCx;
-use crate::intrinsic::llvm;
 
 impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
     pub fn get_or_insert_global(
@@ -100,18 +99,14 @@ pub fn declare_entry_fn(
         let return_type = self.type_i32();
         let variadic = false;
         self.linkage.set(FunctionType::Exported);
-        let func = declare_raw_fn(
+        declare_raw_fn(
             self,
             name,
             callconv,
             return_type,
             &[self.type_i32(), const_string],
             variadic,
-        );
-        // NOTE: it is needed to set the current_func here as well, because get_fn() is not called
-        // for the main function.
-        *self.current_func.borrow_mut() = Some(func);
-        func
+        )
     }
 
     pub fn declare_fn(&self, name: &str, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Function<'gcc> {
@@ -166,19 +161,6 @@ fn declare_raw_fn<'gcc>(
     param_types: &[Type<'gcc>],
     variadic: bool,
 ) -> Function<'gcc> {
-    if name.starts_with("llvm.") {
-        let intrinsic = match name {
-            "llvm.fma.f16" => {
-                // fma is not a target builtin, but a normal builtin, so we handle it differently
-                // here.
-                cx.context.get_builtin_function("fma")
-            }
-            _ => llvm::intrinsic(name, cx),
-        };
-
-        cx.intrinsics.borrow_mut().insert(name.to_string(), intrinsic);
-        return intrinsic;
-    }
     let func = if cx.functions.borrow().contains_key(name) {
         cx.functions.borrow()[name]
     } else {

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -9,7 +9,7 @@
 use gccjit::{ComparisonOp, Function, FunctionType, RValue, ToRValue, UnaryOp};
 #[cfg(feature = "master")]
 use rustc_abi::ExternAbi;
-use rustc_abi::{BackendRepr, HasDataLayout};
+use rustc_abi::{BackendRepr, HasDataLayout, WrappingRange};
 use rustc_codegen_ssa::MemFlags;
 use rustc_codegen_ssa::base::wants_msvc_seh;
 use rustc_codegen_ssa::common::IntPredicate;
@@ -20,19 +20,15 @@
 use rustc_codegen_ssa::traits::MiscCodegenMethods;
 use rustc_codegen_ssa::traits::{
     ArgAbiBuilderMethods, BaseTypeCodegenMethods, BuilderMethods, ConstCodegenMethods,
-    IntrinsicCallBuilderMethods,
+    IntrinsicCallBuilderMethods, LayoutTypeCodegenMethods,
 };
 use rustc_middle::bug;
-#[cfg(feature = "master")]
-use rustc_middle::ty::layout::FnAbiOf;
-use rustc_middle::ty::layout::LayoutOf;
+use rustc_middle::ty::layout::{FnAbiOf, LayoutOf};
 use rustc_middle::ty::{self, Instance, Ty};
 use rustc_span::{Span, Symbol, sym};
 use rustc_target::callconv::{ArgAbi, PassMode};
 
-#[cfg(feature = "master")]
-use crate::abi::FnAbiGccExt;
-use crate::abi::GccType;
+use crate::abi::{FnAbiGccExt, GccType};
 use crate::builder::Builder;
 use crate::common::{SignType, TypeReflection};
 use crate::context::CodegenCx;
@@ -609,6 +605,94 @@ fn codegen_intrinsic_call(
         Ok(())
     }
 
+    fn codegen_llvm_intrinsic_call(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+        args: &[OperandRef<'tcx, Self::Value>],
+        is_cleanup: bool,
+    ) -> Self::Value {
+        let func = if let Some(&func) = self.intrinsic_instances.borrow().get(&instance) {
+            func
+        } else {
+            let sym = self.tcx.symbol_name(instance).name;
+
+            let func = if let Some(func) = self.intrinsics.borrow().get(sym) {
+                *func
+            } else {
+                self.linkage.set(FunctionType::Extern);
+                let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
+                let fn_ty = fn_abi.gcc_type(self);
+
+                let func = match sym {
+                    "llvm.fma.f16" => {
+                        // fma is not a target builtin, but a normal builtin, so we handle it differently
+                        // here.
+                        self.context.get_builtin_function("fma")
+                    }
+                    _ => llvm::intrinsic(sym, self),
+                };
+
+                self.intrinsics.borrow_mut().insert(sym.to_string(), func);
+
+                self.on_stack_function_params
+                    .borrow_mut()
+                    .insert(func, fn_ty.on_stack_param_indices);
+                #[cfg(feature = "master")]
+                for fn_attr in fn_ty.fn_attributes {
+                    func.add_attribute(fn_attr);
+                }
+
+                crate::attributes::from_fn_attrs(self, func, instance);
+
+                func
+            };
+
+            self.intrinsic_instances.borrow_mut().insert(instance, func);
+
+            func
+        };
+        let fn_ptr = func.get_address(None);
+        let fn_ty = fn_ptr.get_type();
+
+        let mut llargs = vec![];
+
+        for arg in args {
+            match arg.val {
+                OperandValue::ZeroSized => {}
+                OperandValue::Immediate(_) => llargs.push(arg.immediate()),
+                OperandValue::Pair(a, b) => {
+                    llargs.push(a);
+                    llargs.push(b);
+                }
+                OperandValue::Ref(op_place_val) => {
+                    let mut llval = op_place_val.llval;
+                    // We can't use `PlaceRef::load` here because the argument
+                    // may have a type we don't treat as immediate, but the ABI
+                    // used for this call is passing it by-value. In that case,
+                    // the load would just produce `OperandValue::Ref` instead
+                    // of the `OperandValue::Immediate` we need for the call.
+                    llval = self.load(self.backend_type(arg.layout), llval, op_place_val.align);
+                    if let BackendRepr::Scalar(scalar) = arg.layout.backend_repr {
+                        if scalar.is_bool() {
+                            self.range_metadata(llval, WrappingRange { start: 0, end: 1 });
+                        }
+                        // We store bools as `i8` so we need to truncate to `i1`.
+                        llval = self.to_immediate_scalar(llval, scalar);
+                    }
+                    llargs.push(llval);
+                }
+            }
+        }
+
+        // FIXME directly use the llvm intrinsic adjustment functions here
+        let llret = self.call(fn_ty, None, None, fn_ptr, &llargs, None, None);
+        if is_cleanup {
+            self.apply_attrs_to_cleanup_callsite(llret);
+        }
+
+        llret
+    }
+
     fn abort(&mut self) {
         let func = self.context.get_builtin_function("abort");
         let func: RValue<'gcc> = unsafe { std::mem::transmute(func) };

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1705,7 +1705,7 @@ pub(crate) fn catch_ret(
         ret.expect("LLVM does not have support for catchret")
     }
 
-    fn check_call<'b>(
+    pub(crate) fn check_call<'b>(
         &mut self,
         typ: &str,
         fn_ty: &'ll Type,

compiler/rustc_codegen_llvm/src/context.rs

@@ -101,6 +101,8 @@ pub(crate) struct FullCx<'ll, 'tcx> {
 
     /// Cache instances of monomorphic and polymorphic items
     pub instances: RefCell<FxHashMap<Instance<'tcx>, &'ll Value>>,
+    /// Cache instances of intrinsics
+    pub intrinsic_instances: RefCell<FxHashMap<Instance<'tcx>, &'ll Value>>,
     /// Cache generated vtables
     pub vtables: RefCell<FxHashMap<(Ty<'tcx>, Option<ty::ExistentialTraitRef<'tcx>>), &'ll Value>>,
     /// Cache of constant strings,
@@ -627,6 +629,7 @@ pub(crate) fn new(
                 tls_model,
                 codegen_unit,
                 instances: Default::default(),
+                intrinsic_instances: Default::default(),
                 vtables: Default::default(),
                 const_str_cache: Default::default(),
                 const_globals: Default::default(),

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1,7 +1,11 @@
 use std::assert_matches::assert_matches;
 use std::cmp::Ordering;
+use std::ffi::c_uint;
+use std::ptr;
 
-use rustc_abi::{Align, BackendRepr, ExternAbi, Float, HasDataLayout, Primitive, Size};
+use rustc_abi::{
+    Align, BackendRepr, ExternAbi, Float, HasDataLayout, Primitive, Size, WrappingRange,
+};
 use rustc_codegen_ssa::base::{compare_simd_types, wants_msvc_seh, wants_wasm_eh};
 use rustc_codegen_ssa::codegen_attrs::autodiff_attrs;
 use rustc_codegen_ssa::common::{IntPredicate, TypeKind};
@@ -28,6 +32,7 @@
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::{gen_call_handling, gen_define_handling};
 use crate::context::CodegenCx;
+use crate::declare::declare_raw_fn;
 use crate::errors::{
     AutoDiffWithoutEnable, AutoDiffWithoutLto, OffloadWithoutEnable, OffloadWithoutFatLTO,
 };
@@ -633,6 +638,99 @@ fn codegen_intrinsic_call(
         Ok(())
     }
 
+    fn codegen_llvm_intrinsic_call(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+        args: &[OperandRef<'tcx, Self::Value>],
+        is_cleanup: bool,
+    ) -> Self::Value {
+        let tcx = self.tcx();
+
+        // FIXME remove usage of fn_abi
+        let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
+        assert!(!fn_abi.ret.is_indirect());
+        let fn_ty = fn_abi.llvm_type(self);
+
+        let fn_ptr = if let Some(&llfn) = self.intrinsic_instances.borrow().get(&instance) {
+            llfn
+        } else {
+            let sym = tcx.symbol_name(instance).name;
+
+            // FIXME use get_intrinsic
+            let llfn = if let Some(llfn) = self.get_declared_value(sym) {
+                llfn
+            } else {
+                // Function addresses in Rust are never significant, allowing functions to
+                // be merged.
+                let llfn = declare_raw_fn(
+                    self,
+                    sym,
+                    fn_abi.llvm_cconv(self),
+                    llvm::UnnamedAddr::Global,
+                    llvm::Visibility::Default,
+                    fn_ty,
+                );
+                fn_abi.apply_attrs_llfn(self, llfn, Some(instance));
+
+                llfn
+            };
+
+            self.intrinsic_instances.borrow_mut().insert(instance, llfn);
+
+            llfn
+        };
+
+        let mut llargs = vec![];
+
+        for arg in args {
+            match arg.val {
+                OperandValue::ZeroSized => {}
+                OperandValue::Immediate(_) => llargs.push(arg.immediate()),
+                OperandValue::Pair(a, b) => {
+                    llargs.push(a);
+                    llargs.push(b);
+                }
+                OperandValue::Ref(op_place_val) => {
+                    let mut llval = op_place_val.llval;
+                    // We can't use `PlaceRef::load` here because the argument
+                    // may have a type we don't treat as immediate, but the ABI
+                    // used for this call is passing it by-value. In that case,
+                    // the load would just produce `OperandValue::Ref` instead
+                    // of the `OperandValue::Immediate` we need for the call.
+                    llval = self.load(self.backend_type(arg.layout), llval, op_place_val.align);
+                    if let BackendRepr::Scalar(scalar) = arg.layout.backend_repr {
+                        if scalar.is_bool() {
+                            self.range_metadata(llval, WrappingRange { start: 0, end: 1 });
+                        }
+                        // We store bools as `i8` so we need to truncate to `i1`.
+                        llval = self.to_immediate_scalar(llval, scalar);
+                    }
+                    llargs.push(llval);
+                }
+            }
+        }
+
+        debug!("call intrinsic {:?} with args ({:?})", instance, llargs);
+        let args = self.check_call("call", fn_ty, fn_ptr, &llargs);
+        let llret = unsafe {
+            llvm::LLVMBuildCallWithOperandBundles(
+                self.llbuilder,
+                fn_ty,
+                fn_ptr,
+                args.as_ptr() as *const &llvm::Value,
+                args.len() as c_uint,
+                ptr::dangling(),
+                0,
+                c"".as_ptr(),
+            )
+        };
+        if is_cleanup {
+            self.apply_attrs_to_cleanup_callsite(llret);
+        }
+
+        llret
+    }
+
     fn abort(&mut self) {
         self.call_intrinsic("llvm.trap", &[], &[]);
     }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -9,12 +9,12 @@
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
 use rustc_middle::ty::print::{with_no_trimmed_paths, with_no_visible_paths};
-use rustc_middle::ty::{self, Instance, Ty};
+use rustc_middle::ty::{self, Instance, Ty, TypeVisitableExt};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::Span;
 use rustc_span::source_map::Spanned;
-use rustc_target::callconv::{ArgAbi, CastTarget, FnAbi, PassMode};
+use rustc_target::callconv::{ArgAbi, ArgAttributes, CastTarget, FnAbi, PassMode};
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
@@ -1036,6 +1036,59 @@ fn codegen_call_terminator(
             _ => bug!("{} is not callable", callee.layout.ty),
         };
 
+        if let Some(instance) = instance
+            && let Some(name) = bx.tcx().codegen_fn_attrs(instance.def_id()).symbol_name
+            && name.as_str().starts_with("llvm.")
+            // This is the only LLVM intrinsic we use that unwinds
+            // FIXME either add unwind support to codegen_llvm_intrinsic_call or replace usage of
+            // this intrinsic with something else
+            && name.as_str() != "llvm.wasm.throw"
+        {
+            assert!(!instance.args.has_infer());
+            assert!(!instance.args.has_escaping_bound_vars());
+
+            let result_layout =
+                self.cx.layout_of(self.monomorphized_place_ty(destination.as_ref()));
+
+            let return_dest = if result_layout.is_zst() {
+                ReturnDest::Nothing
+            } else if let Some(index) = destination.as_local() {
+                match self.locals[index] {
+                    LocalRef::Place(dest) => ReturnDest::Store(dest),
+                    LocalRef::UnsizedPlace(_) => bug!("return type must be sized"),
+                    LocalRef::PendingOperand => {
+                        // Handle temporary places, specifically `Operand` ones, as
+                        // they don't have `alloca`s.
+                        ReturnDest::DirectOperand(index)
+                    }
+                    LocalRef::Operand(_) => bug!("place local already assigned to"),
+                }
+            } else {
+                ReturnDest::Store(self.codegen_place(bx, destination.as_ref()))
+            };
+
+            let args =
+                args.into_iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect::<Vec<_>>();
+
+            self.set_debug_loc(bx, source_info);
+
+            let llret =
+                bx.codegen_llvm_intrinsic_call(instance, &args, self.mir[helper.bb].is_cleanup);
+
+            if let Some(target) = target {
+                self.store_return(
+                    bx,
+                    return_dest,
+                    &ArgAbi { layout: result_layout, mode: PassMode::Direct(ArgAttributes::new()) },
+                    llret,
+                );
+                return helper.funclet_br(self, bx, target, mergeable_succ);
+            } else {
+                bx.unreachable();
+                return MergingSucc::False;
+            }
+        }
+
         // FIXME(eddyb) avoid computing this if possible, when `instance` is
         // available - right now `sig` is only needed for getting the `abi`
         // and figuring out how many extra args were passed to a C-variadic `fn`.

compiler/rustc_codegen_ssa/src/traits/intrinsic.rs

@@ -25,6 +25,13 @@ fn codegen_intrinsic_call(
         span: Span,
     ) -> Result<(), ty::Instance<'tcx>>;
 
+    fn codegen_llvm_intrinsic_call(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+        args: &[OperandRef<'tcx, Self::Value>],
+        is_cleanup: bool,
+    ) -> Self::Value;
+
     fn abort(&mut self);
     fn assume(&mut self, val: Self::Value);
     fn expect(&mut self, cond: Self::Value, expected: bool) -> Self::Value;