Add target_feature_available_at_call_site intrinsic

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -20,6 +20,8 @@ mod simd;
 use cranelift_codegen::ir::{
     AtomicRmwOp, BlockArg, ExceptionTableData, ExceptionTableItem, ExceptionTag,
 };
+use rustc_codegen_ssa::target_features;
+use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_middle::ty;
 use rustc_middle::ty::GenericArgsRef;
 use rustc_middle::ty::layout::ValidityRequirement;
@@ -834,6 +836,29 @@ fn codegen_regular_intrinsic_call<'tcx>(
             ret.write_cvalue(fx, caller_location);
         }
 
+        sym::target_feature_available_at_call_site => {
+            intrinsic_args!(fx, args => (); intrinsic);
+
+            let enabled = match target_features::target_feature_available_at_call_site_intrinsic(
+                fx.tcx,
+                source_info.span,
+                fx.instance.def_id(),
+                generic_args,
+            ) {
+                target_features::TargetFeatureAvailableAtCallSite::Known(enabled) => enabled,
+                // SSA already handles the easy case where the caller function has the feature enabled.
+                // Cranelift doesn't (yet) have the equivalent LLVM pass to replace a marker post-inlining,
+                // so report that the feature is unavailable at the call site.
+                target_features::TargetFeatureAvailableAtCallSite::CheckBackend(_) => false,
+            };
+
+            let ret_val = CValue::by_val(
+                fx.bcx.ins().iconst(types::I8, i64::from(enabled)),
+                fx.layout_of(fx.tcx.types.bool),
+            );
+            ret.write_cvalue(fx, ret_val);
+        }
+
         _ if intrinsic.as_str().starts_with("atomic_fence") => {
             intrinsic_args!(fx, args => (); intrinsic);
 

compiler/rustc_codegen_gcc/src/intrinsic/mod.rs

@@ -24,7 +24,7 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_middle::ty::layout::FnAbiOf;
 use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::{self, Instance, Ty};
-use rustc_middle::{bug, span_bug};
+use rustc_middle::{bug, mir, span_bug};
 use rustc_span::{Span, Symbol, sym};
 use rustc_target::callconv::{ArgAbi, PassMode};
 
@@ -195,6 +195,7 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'gcc, 'tc
         &mut self,
         instance: Instance<'tcx>,
         args: &[OperandRef<'tcx, RValue<'gcc>>],
+        _: &[mir::Operand<'tcx>],
         result_layout: ty::layout::TyAndLayout<'tcx>,
         result_place: Option<PlaceValue<RValue<'gcc>>>,
         span: Span,
@@ -594,6 +595,16 @@ impl<'a, 'gcc, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'gcc, 'tc
         IntrinsicResult::WroteIntoPlace
     }
 
+    fn codegen_target_feature_available_at_call_site(
+        &mut self,
+        _rust_feature_name: &str,
+    ) -> RValue<'gcc> {
+        // SSA already handles the easy case where the caller function has the feature enabled.
+        // GCC doesn't (yet) have the equivalent LLVM pass to replace a marker post-inlining,
+        // so report that the feature is unavailable at the call site.
+        self.const_bool(false)
+    }
+
     fn codegen_llvm_intrinsic_call(
         &mut self,
         instance: ty::Instance<'tcx>,

compiler/rustc_codegen_llvm/src/context.rs

@@ -36,7 +36,7 @@ use crate::back::write::to_llvm_code_model;
 use crate::builder::gpu_offload::{OffloadGlobals, OffloadKernelGlobals};
 use crate::callee::get_fn;
 use crate::debuginfo::metadata::apply_vcall_visibility_metadata;
-use crate::llvm::{self, Metadata, MetadataKindId, Module, Type, Value};
+use crate::llvm::{self, Metadata, MetadataKindId, Module, TargetMachine, Type, Value};
 use crate::{attributes, common, coverageinfo, debuginfo, llvm_util};
 
 /// `TyCtxt` (and related cache datastructures) can't be move between threads.
@@ -92,6 +92,7 @@ pub(crate) type CodegenCx<'ll, 'tcx> = GenericCx<'ll, FullCx<'ll, 'tcx>>;
 pub(crate) struct FullCx<'ll, 'tcx> {
     pub tcx: TyCtxt<'tcx>,
     pub scx: SimpleCx<'ll>,
+    pub tm: &'ll TargetMachine,
     pub use_dll_storage_attrs: bool,
     pub tls_model: llvm::ThreadLocalMode,
 
@@ -660,6 +661,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             FullCx {
                 tcx,
                 scx: SimpleCx::new(llmod, llcx, tcx.data_layout.pointer_size()),
+                tm: llvm_module.tm.raw(),
                 use_dll_storage_attrs,
                 tls_model,
                 codegen_unit,

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -17,7 +17,7 @@ use rustc_codegen_ssa::traits::*;
 use rustc_hir as hir;
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_hir::find_attr;
-use rustc_middle::mir::BinOp;
+use rustc_middle::mir::{self, BinOp};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv, LayoutOf};
 use rustc_middle::ty::offload_meta::OffloadMetadata;
 use rustc_middle::ty::{self, GenericArgsRef, Instance, SimdAlign, Ty, TyCtxt, TypingEnv};
@@ -32,6 +32,7 @@ use rustc_target::spec::{Arch, LlvmAbi};
 use tracing::debug;
 
 use crate::abi::FnAbiLlvmExt;
+use crate::attributes;
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::{
@@ -177,6 +178,7 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
         &mut self,
         instance: ty::Instance<'tcx>,
         args: &[OperandRef<'tcx, &'ll Value>],
+        _: &[mir::Operand<'tcx>],
         result_layout: ty::layout::TyAndLayout<'tcx>,
         result_place: Option<PlaceValue<&'ll Value>>,
         span: Span,
@@ -923,6 +925,47 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
         }
     }
 
+    fn codegen_target_feature_available_at_call_site(
+        &mut self,
+        rust_feature_name: &str,
+    ) -> &'ll Value {
+        // SSA already returned `true` when this Rust feature is enabled on the current
+        // function. If there is no LLVM mapping, we cannot do any better from here.
+        let Some(llvm_feature) = crate::llvm_util::to_llvm_features(self.sess(), rust_feature_name)
+        else {
+            return self.const_bool(false);
+        };
+
+        // Target features can't be lost by inlining, so early return if it's already present.
+        // SSA should have already done this, but this is most accurate for implied features.
+        let mut llvm_features = llvm_feature.into_iter();
+        let llvm_feature_name = llvm_features.next().unwrap().to_string();
+        let enabled = llvm::FunctionHasTargetFeature(self.cx.tm, self.llfn(), &llvm_feature_name);
+        if enabled {
+            return self.const_bool(true);
+        }
+
+        // If we're not optimizing, we won't run the LLVM pass
+        if self.sess().opts.optimize == rustc_session::config::OptLevel::No {
+            return self.const_bool(false);
+        }
+
+        // Generate the marker to be replaced by the LLVM pass
+        let fn_ty = self.type_func(&[], self.type_i1());
+        let llfn = self.cx.declare_cfn(
+            "rust.target_feature_available_at_call_site",
+            llvm::UnnamedAddr::No,
+            fn_ty,
+        );
+        let nounwind = llvm::AttributeKind::NoUnwind.create_attr(self.cx.llcx);
+        attributes::apply_to_llfn(llfn, llvm::AttributePlace::Function, &[nounwind]);
+        let call = self.call(fn_ty, None, None, llfn, &[], None, None);
+        let kind = self.cx.get_md_kind_id("rust.target_feature");
+        let feature = self.cx.create_metadata(llvm_feature_name.as_bytes());
+        self.cx.set_metadata_node(call, kind, &[feature]);
+        call
+    }
+
     fn codegen_llvm_intrinsic_call(
         &mut self,
         instance: ty::Instance<'tcx>,

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -944,7 +944,6 @@ unsafe extern "C" {
     pub(crate) fn LLVMFloatTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMDoubleTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMFP128TypeInContext(C: &Context) -> &Type;
-
     // Operations on non-IEEE real types
     pub(crate) fn LLVMBFloatTypeInContext(C: &Context) -> &Type;
 
@@ -2406,6 +2405,12 @@ unsafe extern "C" {
 
     pub(crate) fn LLVMRustHasFeature(T: &TargetMachine, s: *const c_char) -> bool;
     pub(crate) fn LLVMRustTargetHasMnemonic(T: &TargetMachine, s: *const c_char) -> bool;
+    pub(crate) fn LLVMRustFunctionHasTargetFeature(
+        TM: &TargetMachine,
+        F: &Value,
+        Feature: *const c_char,
+        FeatureLen: size_t,
+    ) -> bool;
 
     pub(crate) fn LLVMRustPrintTargetCPUs(TM: &TargetMachine, OutStr: &RustString);
     pub(crate) fn LLVMRustGetTargetFeaturesCount(T: &TargetMachine) -> size_t;

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -51,6 +51,14 @@ pub(crate) fn RemoveStringAttrFromFn<'ll>(llfn: &'ll Value, name: &str) {
     unsafe { LLVMRustRemoveFnAttribute(llfn, name.as_c_char_ptr(), name.len()) }
 }
 
+pub(crate) fn FunctionHasTargetFeature<'ll>(
+    tm: &'ll TargetMachine,
+    llfn: &'ll Value,
+    feature: &str,
+) -> bool {
+    unsafe { LLVMRustFunctionHasTargetFeature(tm, llfn, feature.as_c_char_ptr(), feature.len()) }
+}
+
 pub(crate) fn AddCallSiteAttributes<'ll>(
     callsite: &'ll Value,
     idx: AttributePlace,

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -1008,13 +1008,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                             span_bug!(self.mir.span, "can't directly store to unaligned value");
                         }
 
-                        let args: Vec<_> =
+                        let mir_args: Vec<_> = args.iter().map(|arg| arg.node.clone()).collect();
+                        let codegen_args: Vec<_> =
                             args.iter().map(|arg| self.codegen_operand(bx, &arg.node)).collect();
 
                         let intrinsic_result = self.codegen_intrinsic_call(
                             bx,
                             instance,
-                            &args,
+                            &codegen_args,
+                            &mir_args,
                             result_layout,
                             result_place,
                             source_info,

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -1,5 +1,5 @@
 use rustc_abi::{Align, FieldIdx, WrappingRange};
-use rustc_middle::mir::SourceInfo;
+use rustc_middle::mir::{self, SourceInfo};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::OptLevel;
@@ -13,7 +13,7 @@ use crate::common::{AtomicRmwBinOp, SynchronizationScope};
 use crate::errors::InvalidMonomorphization;
 use crate::mir::operand::OperandRefBuilder;
 use crate::traits::*;
-use crate::{MemFlags, meth, size_of_val};
+use crate::{MemFlags, meth, size_of_val, target_features};
 
 fn copy_intrinsic<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     bx: &mut Bx,
@@ -59,6 +59,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx: &mut Bx,
         instance: ty::Instance<'tcx>,
         args: &[OperandRef<'tcx, Bx::Value>],
+        mir_args: &[mir::Operand<'tcx>],
         result_layout: ty::layout::TyAndLayout<'tcx>,
         result_place: Option<PlaceValue<Bx::Value>>,
         source_info: SourceInfo,
@@ -124,6 +125,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 | sym::atomic_fence
                 | sym::atomic_singlethreadfence
                 | sym::caller_location
+                | sym::target_feature_available_at_call_site
                 | sym::return_address => {}
                 _ => {
                     span_bug!(
@@ -589,10 +591,34 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 OperandValue::ZeroSized
             }
 
+            sym::target_feature_available_at_call_site => {
+                match target_features::target_feature_available_at_call_site_intrinsic(
+                    bx.tcx(),
+                    span,
+                    self.instance.def_id(),
+                    fn_args,
+                ) {
+                    target_features::TargetFeatureAvailableAtCallSite::Known(enabled) => {
+                        OperandValue::Immediate(bx.const_bool(enabled))
+                    }
+                    target_features::TargetFeatureAvailableAtCallSite::CheckBackend(feature) => {
+                        OperandValue::Immediate(
+                            bx.codegen_target_feature_available_at_call_site(feature.as_str()),
+                        )
+                    }
+                }
+            }
+
             _ => {
                 // Need to use backend-specific things in the implementation.
-                let result =
-                    bx.codegen_intrinsic_call(instance, args, result_layout, result_place, span);
+                let result = bx.codegen_intrinsic_call(
+                    instance,
+                    args,
+                    mir_args,
+                    result_layout,
+                    result_place,
+                    span,
+                );
                 if let IntrinsicResult::Operand(op) = result {
                     op
                 } else {

compiler/rustc_codegen_ssa/src/target_features.rs

@@ -5,7 +5,8 @@ use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE, LocalDefId};
 use rustc_middle::middle::codegen_fn_attrs::{TargetFeature, TargetFeatureKind};
 use rustc_middle::query::Providers;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::span_bug;
+use rustc_middle::ty::{self, TyCtxt};
 use rustc_session::Session;
 use rustc_session::errors::feature_err;
 use rustc_session::lint::builtin::AARCH64_SOFTFLOAT_NEON;
@@ -130,7 +131,7 @@ pub(crate) fn from_target_feature_attr(
 
 /// Computes the set of target features used in a function for the purposes of
 /// inline assembly.
-fn asm_target_features(tcx: TyCtxt<'_>, did: DefId) -> &FxIndexSet<Symbol> {
+pub(crate) fn asm_target_features(tcx: TyCtxt<'_>, did: DefId) -> &FxIndexSet<Symbol> {
     let mut target_features = tcx.sess.unstable_target_features.clone();
     if tcx.def_kind(did).has_codegen_attrs() {
         let attrs = tcx.codegen_fn_attrs(did);
@@ -569,3 +570,64 @@ pub(crate) fn provide(providers: &mut Providers) {
         ..*providers
     }
 }
+
+#[derive(Copy, Clone, Debug)]
+pub enum TargetFeatureAvailableAtCallSite {
+    /// The feature's presence is known immediately, including `false` for error conditions.
+    Known(bool),
+    /// It's not known if the feature is available at the call site. It's safe to return `false`,
+    /// but the backend may be able to refine the answer (e.g. the LLVM post-inline pass)
+    CheckBackend(Symbol),
+}
+
+/// Implementation for the `target_feature_available_at_call_site` intrinsic.
+///
+/// This function performs the backend-independent implementation:
+/// * If the feature is enabled for the current function, we can immediately return true.
+/// * If the feature is not enabled for the current function, defer to the backend to account
+///   for inlining.
+/// * If an error occurs, emit the diagnostic and return false.
+///
+/// When deferred to the backend, it's always safe to return false, but optimization opportunities
+/// are lost.
+/// The LLVM backend, for example, emits a marker function that is replaced with true or false
+/// with a post-inlining pass.
+pub fn target_feature_available_at_call_site_intrinsic<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    span: Span,
+    did: DefId,
+    generic_args: ty::GenericArgsRef<'tcx>,
+) -> TargetFeatureAvailableAtCallSite {
+    let Some(feature_bytes) =
+        generic_args.const_at(0).try_to_value().and_then(|feature| feature.try_to_raw_bytes(tcx))
+    else {
+        span_bug!(span, "wrong const argument for target_feature_available_at_call_site");
+    };
+
+    let feature_len =
+        feature_bytes.iter().position(|&byte| byte == 0).unwrap_or(feature_bytes.len());
+    let Some(feature_name) = std::str::from_utf8(&feature_bytes[..feature_len]).ok() else {
+        tcx.dcx().span_err(
+            span,
+            "`target_feature_available_at_call_site` requires a UTF-8 feature name",
+        );
+        return TargetFeatureAvailableAtCallSite::Known(false);
+    };
+    let feature = Symbol::intern(feature_name);
+
+    let rust_target_features = tcx.rust_target_features(LOCAL_CRATE);
+    let Some(&stability) = rust_target_features.get(feature_name) else {
+        tcx.dcx().span_err(span, format!("unknown target feature `{feature}`"));
+        return TargetFeatureAvailableAtCallSite::Known(false);
+    };
+    if let Err(reason) = stability.toggle_allowed() {
+        tcx.dcx().span_err(span, format!("cannot use target feature `{feature}`: {reason}"));
+        return TargetFeatureAvailableAtCallSite::Known(false);
+    }
+
+    if asm_target_features(tcx, did).contains(&feature) {
+        TargetFeatureAvailableAtCallSite::Known(true)
+    } else {
+        TargetFeatureAvailableAtCallSite::CheckBackend(feature)
+    }
+}