Prototype loop unrolling hint attributes

compiler/rustc_attr_parsing/src/attributes/loop_match.rs

@@ -3,7 +3,7 @@ use super::prelude::*;
 pub(crate) struct LoopMatchParser;
 impl NoArgsAttributeParser for LoopMatchParser {
     const PATH: &[Symbol] = &[sym::loop_match];
-    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[Allow(Target::Expression)]);
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[Allow(Target::Loop)]);
     const CREATE: fn(Span) -> AttributeKind = AttributeKind::LoopMatch;
 }
 

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -68,6 +68,7 @@ pub(crate) mod stability;
 pub(crate) mod test_attrs;
 pub(crate) mod traits;
 pub(crate) mod transparency;
+pub(crate) mod unroll;
 pub(crate) mod util;
 
 type AcceptFn<T> = for<'sess> fn(&mut T, &mut AcceptContext<'_, 'sess>, &ArgParser);

compiler/rustc_attr_parsing/src/attributes/unroll.rs

@@ -0,0 +1,56 @@
+use rustc_ast::{LitIntType, LitKind};
+use rustc_hir::attrs::UnrollAttr;
+use rustc_lint_defs::builtin::ILL_FORMED_ATTRIBUTE_INPUT;
+
+use super::prelude::*;
+
+pub(crate) struct UnrollParser;
+impl SingleAttributeParser for UnrollParser {
+    const PATH: &[Symbol] = &[sym::unroll];
+    const ALLOWED_TARGETS: AllowedTargets = AllowedTargets::AllowList(&[
+        Allow(Target::Loop),
+        Allow(Target::ForLoop),
+        Allow(Target::While),
+    ]);
+    const TEMPLATE: AttributeTemplate = template!(
+        Word,
+        List: &["always", "never", "<integer>"],
+        "https://doc.rust-lang.org/reference/attributes/codegen.html#the-inline-attribute"
+    );
+
+    fn convert(cx: &mut AcceptContext<'_, '_>, args: &ArgParser) -> Option<AttributeKind> {
+        match args {
+            ArgParser::NoArgs => Some(AttributeKind::Unroll(UnrollAttr::Hint)),
+            ArgParser::List(list) => {
+                let l = cx.expect_single(list)?;
+
+                if let Some(lit) = l.as_lit() {
+                    if let LitKind::Int(val, LitIntType::Unsuffixed) = lit.kind {
+                        let Ok(val) = u32::try_from(val.get()) else {
+                            cx.adcx().expected_integer_literal_in_range(
+                                l.span(),
+                                0,
+                                u32::MAX as isize,
+                            );
+                            return None;
+                        };
+                        return Some(AttributeKind::Unroll(UnrollAttr::Count(val)));
+                    }
+                }
+
+                match l.meta_item().and_then(|i| i.path().word_sym()) {
+                    Some(sym::always) => Some(AttributeKind::Unroll(UnrollAttr::Always)),
+                    Some(sym::never) => Some(AttributeKind::Unroll(UnrollAttr::Never)),
+                    _ => {
+                        cx.adcx().expected_specific_argument(l.span(), &[sym::always, sym::never]);
+                        None
+                    }
+                }
+            }
+            ArgParser::NameValue(_) => {
+                cx.adcx().warn_ill_formed_attribute_input(ILL_FORMED_ATTRIBUTE_INPUT);
+                return None;
+            }
+        }
+    }
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -58,6 +58,7 @@ use crate::attributes::stability::*;
 use crate::attributes::test_attrs::*;
 use crate::attributes::traits::*;
 use crate::attributes::transparency::*;
+use crate::attributes::unroll::*;
 use crate::attributes::{AttributeParser as _, AttributeSafety, Combine, Single, WithoutArgs};
 use crate::parser::{
     ArgParser, MetaItemListParser, MetaItemOrLitParser, MetaItemParser, NameValueParser,
@@ -224,6 +225,7 @@ attribute_parsers!(
         Single<ShouldPanicParser>,
         Single<TestRunnerParser>,
         Single<TypeLengthLimitParser>,
+        Single<UnrollParser>,
         Single<WindowsSubsystemParser>,
         Single<WithoutArgs<AllowInternalUnsafeParser>>,
         Single<WithoutArgs<AutomaticallyDerivedParser>>,

compiler/rustc_attr_parsing/src/target_checking.rs

@@ -465,5 +465,8 @@ pub(crate) const ALL_TARGETS: &'static [Policy] = {
             kind: rustc_hir::target::GenericParamKind::Type,
             has_default: true,
         }),
+        Allow(Target::Loop),
+        Allow(Target::ForLoop),
+        Allow(Target::While),
     ]
 };

compiler/rustc_codegen_llvm/src/builder.rs

@@ -14,6 +14,8 @@ use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::small_c_str::SmallCStr;
+use rustc_hir::Attribute;
+use rustc_hir::attrs::{AttributeKind, UnrollAttr};
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::ty::layout::{
@@ -336,6 +338,50 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn br_with_attrs(&mut self, dest: &'ll BasicBlock, attributes: &[Attribute]) {
+        unsafe {
+            let val = llvm::LLVMBuildBr(self.llbuilder, dest);
+
+            let mut nodes = Vec::new();
+
+            for attribute in attributes {
+                if let Attribute::Parsed(AttributeKind::Unroll(unroll)) = attribute {
+                    // UnrollAttr::Count needs a second operand, the provided count, but the other
+                    // unroll hints do not.
+                    let md_node = if let UnrollAttr::Count(count) = unroll {
+                        let unroll_meta = self.create_metadata("llvm.loop.unroll.count".as_bytes());
+                        let count =
+                            llvm::LLVMValueAsMetadata(self.get_const_i32(u64::from(*count)));
+                        self.md_node_in_context(&[unroll_meta, count])
+                    } else {
+                        let metadata_str = match unroll {
+                            UnrollAttr::Hint => "llvm.loop.unroll.enable",
+                            UnrollAttr::Always => "llvm.loop.unroll.full",
+                            UnrollAttr::Never => "llvm.loop.unroll.disable",
+                            _ => unreachable!(),
+                        };
+                        let unroll_meta = self.create_metadata(metadata_str.as_bytes());
+                        self.md_node_in_context(&[unroll_meta])
+                    };
+                    nodes.push(md_node);
+                }
+            }
+
+            if nodes.len() > 0 {
+                nodes.insert(0, nodes[0]);
+                // Create the loop metadata node
+                let loop_meta_mdnode = self.set_metadata_node(val, llvm::MD_loop, &nodes);
+
+                // Look up the metadata node as a value
+                let loop_meta_val = llvm::LLVMGetMetadata(val, llvm::MD_loop).unwrap();
+
+                // Replace the first entry with a reference to itself
+                // This is required by LLVM. See the LangRef page for llvm.loop metadata.
+                llvm::LLVMReplaceMDNodeOperandWith(loop_meta_val, 0, loop_meta_mdnode);
+            }
+        }
+    }
+
     fn cond_br(
         &mut self,
         cond: &'ll Value,

compiler/rustc_codegen_llvm/src/context.rs

@@ -1073,9 +1073,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         instruction: &'ll Value,
         kind_id: MetadataKindId,
         md_list: &[&'ll Metadata],
-    ) {
+    ) -> &'ll Metadata {
         let md = self.md_node_in_context(md_list);
         self.set_metadata(instruction, kind_id, md);
+        md
     }
 
     /// Helper method for the sequence of calls:

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -961,6 +961,10 @@ unsafe extern "C" {
     pub(crate) fn LLVMGetValueName2(Val: &Value, Length: *mut size_t) -> *const c_char;
     pub(crate) fn LLVMSetValueName2(Val: &Value, Name: *const c_char, NameLen: size_t);
     pub(crate) fn LLVMReplaceAllUsesWith<'a>(OldVal: &'a Value, NewVal: &'a Value);
+    pub(crate) safe fn LLVMGetMetadata<'a>(
+        Val: &'a Value,
+        KindID: MetadataKindId,
+    ) -> Option<&'a Value>;
     pub(crate) safe fn LLVMSetMetadata<'a>(Val: &'a Value, KindID: MetadataKindId, Node: &'a Value);
     pub(crate) fn LLVMGlobalSetMetadata<'a>(
         Val: &'a Value,
@@ -990,6 +994,7 @@ unsafe extern "C" {
         Name: *const c_char,
         Val: &'a Value,
     );
+    pub(crate) fn LLVMReplaceMDNodeOperandWith(Val: &Value, index: u32, replacement: &Metadata);
 
     // Operations on scalar constants
     pub(crate) fn LLVMConstInt(IntTy: &Type, N: c_ulonglong, SignExtend: Bool) -> &Value;

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -4,6 +4,7 @@ use rustc_abi::{Align, BackendRepr, ExternAbi, HasDataLayout, Reg, Size, Wrappin
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_data_structures::packed::Pu128;
+use rustc_hir::Attribute;
 use rustc_hir::lang_items::LangItem;
 use rustc_lint_defs::builtin::TAIL_CALL_TRACK_CALLER;
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
@@ -135,6 +136,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
         bx: &mut Bx,
         target: mir::BasicBlock,
         mergeable_succ: bool,
+        attributes: &[Attribute],
     ) -> MergingSucc {
         let (needs_landing_pad, is_cleanupret) = self.llbb_characteristics(fx, target);
         if mergeable_succ && !needs_landing_pad && !is_cleanupret {
@@ -150,7 +152,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                 // to a trampoline.
                 bx.cleanup_ret(self.funclet(fx).unwrap(), Some(lltarget));
             } else {
-                bx.br(lltarget);
+                bx.br_with_attrs(lltarget, attributes);
             }
             MergingSucc::False
         }
@@ -285,7 +287,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
                     bx.lifetime_end(tmp, size);
                 }
                 fx.store_return(bx, ret_dest, &fn_abi.ret, llret);
-                self.funclet_br(fx, bx, target, mergeable_succ)
+                self.funclet_br(fx, bx, target, mergeable_succ, &[])
             } else {
                 bx.unreachable();
                 MergingSucc::False
@@ -353,7 +355,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
             bx.codegen_inline_asm(template, operands, options, line_spans, instance, None, None);
 
             if let Some(target) = destination {
-                self.funclet_br(fx, bx, target, mergeable_succ)
+                self.funclet_br(fx, bx, target, mergeable_succ, &[])
             } else {
                 bx.unreachable();
                 MergingSucc::False
@@ -608,7 +610,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
         if let ty::InstanceKind::DropGlue(_, None) = drop_fn.def {
             // we don't actually need to drop anything.
-            return helper.funclet_br(self, bx, target, mergeable_succ);
+            return helper.funclet_br(self, bx, target, mergeable_succ, &[]);
         }
 
         let place = self.codegen_place(bx, location.as_ref());
@@ -717,7 +719,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
         // Don't codegen the panic block if success if known.
         if const_cond == Some(expected) {
-            return helper.funclet_br(self, bx, target, mergeable_succ);
+            return helper.funclet_br(self, bx, target, mergeable_succ, &[]);
         }
 
         // Because we're branching to a panic block (either a `#[cold]` one
@@ -851,7 +853,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         if is_valid {
             // a NOP
             let target = target.unwrap();
-            return Some(helper.funclet_br(self, bx, target, mergeable_succ));
+            return Some(helper.funclet_br(self, bx, target, mergeable_succ, &[]));
         }
 
         let layout = bx.layout_of(ty);
@@ -925,7 +927,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     ty::InstanceKind::DropGlue(_, None) => {
                         // Empty drop glue; a no-op.
                         let target = target.unwrap();
-                        return helper.funclet_br(self, bx, target, mergeable_succ);
+                        return helper.funclet_br(self, bx, target, mergeable_succ, &[]);
                     }
                     ty::InstanceKind::Intrinsic(def_id) => {
                         let intrinsic = bx.tcx().intrinsic(def_id).unwrap();
@@ -1013,7 +1015,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         match intrinsic_result {
                             IntrinsicResult::Operand(_) | IntrinsicResult::WroteIntoPlace => {
                                 return if let Some(target) = target {
-                                    helper.funclet_br(self, bx, target, mergeable_succ)
+                                    helper.funclet_br(self, bx, target, mergeable_succ, &[])
                                 } else {
                                     bx.unreachable();
                                     MergingSucc::False
@@ -1123,7 +1125,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     &ArgAbi { layout: result_layout, mode: PassMode::Direct(ArgAttributes::new()) },
                     llret,
                 );
-                return helper.funclet_br(self, bx, target, mergeable_succ);
+                return helper.funclet_br(self, bx, target, mergeable_succ, &[]);
             } else {
                 bx.unreachable();
                 return MergingSucc::False;
@@ -1567,7 +1569,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             mir::TerminatorKind::Goto { target } => {
-                helper.funclet_br(self, bx, target, mergeable_succ())
+                helper.funclet_br(self, bx, target, mergeable_succ(), &terminator.attributes)
             }
 
             mir::TerminatorKind::SwitchInt { ref discr, ref targets } => {

compiler/rustc_codegen_ssa/src/traits/builder.rs

@@ -2,6 +2,7 @@ use std::assert_matches;
 use std::ops::Deref;
 
 use rustc_abi::{Align, Scalar, Size, WrappingRange};
+use rustc_hir::Attribute;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::mir;
 use rustc_middle::ty::layout::{FnAbiOf, LayoutOf, TyAndLayout};
@@ -77,6 +78,9 @@ pub trait BuilderMethods<'a, 'tcx>:
     fn ret_void(&mut self);
     fn ret(&mut self, v: Self::Value);
     fn br(&mut self, dest: Self::BasicBlock);
+    fn br_with_attrs(&mut self, dest: Self::BasicBlock, _attributes: &[Attribute]) {
+        self.br(dest)
+    }
     fn cond_br(
         &mut self,
         cond: Self::Value,

compiler/rustc_feature/src/builtin_attrs.rs

@@ -401,6 +401,8 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
     // - https://github.com/rust-lang/rust/issues/130494
     gated!(pin_v2, pin_ergonomics, experimental!(pin_v2)),
 
+    gated!(unroll, loop_hints, experimental!(loop_hints)),
+
     // ==========================================================================
     // Internal attributes: Stability, deprecation, and unsafe:
     // ==========================================================================

compiler/rustc_feature/src/unstable.rs

@@ -606,6 +606,8 @@ declare_features! (
     (unstable, link_arg_attribute, "1.76.0", Some(99427)),
     /// Target features on loongarch.
     (unstable, loongarch_target_feature, "1.73.0", Some(150252)),
+    /// Allows use of loop optimization hints via attributes.
+    (unstable, loop_hints, "CURRENT_RUSTC_VERSION", Some(156874)),
     /// Allows fused `loop`/`match` for direct intraprocedural jumps.
     (incomplete, loop_match, "1.90.0", Some(132306)),
     /// Target features on m68k.

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -129,6 +129,14 @@ impl InlineAttr {
     }
 }
 
+#[derive(Copy, Clone, PartialEq, Encodable, Decodable, Debug, StableHash, PrintAttribute)]
+pub enum UnrollAttr {
+    Hint,
+    Always,
+    Never,
+    Count(u32),
+}
+
 #[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq, StableHash, PrintAttribute)]
 pub enum InstructionSetAttr {
     ArmA32,
@@ -1615,6 +1623,9 @@ pub enum AttributeKind {
         limit: Limit,
     },
 
+    /// Represents `#[unroll]`
+    Unroll(UnrollAttr),
+
     /// Represents `#[unstable_feature_bound]`.
     UnstableFeatureBound(ThinVec<(Symbol, Span)>),
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -197,6 +197,7 @@ impl AttributeKind {
             ThreadLocal => No,
             TrackCaller(..) => Yes,
             TypeLengthLimit { .. } => No,
+            Unroll(..) => No,
             UnstableFeatureBound(..) => No,
             UnstableRemoved(..) => Yes,
             Used { .. } => No,

compiler/rustc_hir/src/attrs/pretty_printing.rs

@@ -192,7 +192,7 @@ macro_rules! print_tup {
 
 print_tup!(A B C D E F G H);
 print_skip!(Span, (), ErrorGuaranteed, AttrId);
-print_disp!(u8, u16, u32, u128, usize, bool, NonZero<u32>, Limit);
+print_disp!(u8, u16, u32, u64, u128, usize, bool, NonZero<u32>, Limit);
 print_debug!(
     Symbol,
     Ident,