Auto merge of #156816 - saethlin:loop-attributes, r=JonathanBrouwer,folkertdev

Add unstable loop unrolling hint attributes



Tracking issue: #156874

This adds as new attribute `#[unroll]`/`#[unroll(full)]`/`#[unroll(never)]`/`#[unroll(16)]` (or any u32).

`#[unroll]` is behind a new feature gate `#![feature(loop_hints)]` because I intend to add an attribute for loop vectorization as well. If a user wants to turn off loop unrolling to locally minimize code size, LLVM may vectorize the loop even though it isn't unrolled which can produce a similar code size explosion.

Author: bors

compiler/rustc_attr_parsing/src/attributes/loop_match.rs

@@ -5,7 +5,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 STABILITY: AttributeStability = unstable!(loop_match);
     const CREATE: fn(Span) -> AttributeKind = AttributeKind::LoopMatch;
 }

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -74,6 +74,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,54 @@
+use rustc_ast::{LitIntType, LitKind};
+use rustc_feature::AttributeStability;
+use rustc_hir::attrs::UnrollAttr;
+
+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 STABILITY: AttributeStability = unstable!(loop_hints);
+    const TEMPLATE: AttributeTemplate = template!(
+        Word,
+        List: &["full", "never", "<integer>"]
+    );
+
+    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()
+                    && let LitKind::Int(val, LitIntType::Unsuffixed) = lit.kind
+                {
+                    if let Ok(val) = u32::try_from(val.get()) {
+                        return Some(AttributeKind::Unroll(UnrollAttr::Count(val)));
+                    } else {
+                        cx.adcx().expected_integer_literal_in_range(l.span(), 0, u32::MAX as isize);
+                        return None;
+                    }
+                }
+
+                match l.meta_item_no_args().and_then(|i| i.path().word_sym()) {
+                    Some(sym::full) => Some(AttributeKind::Unroll(UnrollAttr::Full)),
+                    Some(sym::never) => Some(AttributeKind::Unroll(UnrollAttr::Never)),
+                    _ => {
+                        cx.adcx().expected_specific_argument(l.span(), &[sym::full, sym::never]);
+                        None
+                    }
+                }
+            }
+            ArgParser::NameValue(_) => {
+                let inner_span = cx.inner_span;
+                cx.adcx().expected_list_or_no_args(inner_span);
+                return None;
+            }
+        }
+    }
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -63,6 +63,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,
@@ -232,6 +233,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

@@ -524,5 +524,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,7 @@ 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::attrs::{AttributeKind, UnrollAttr};
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::ty::layout::{
@@ -336,6 +337,51 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         }
     }
 
+    fn br_with_attrs(&mut self, dest: &'ll BasicBlock, attributes: &[AttributeKind]) {
+        unsafe {
+            let val = llvm::LLVMBuildBr(self.llbuilder, dest);
+
+            let mut nodes = Vec::new();
+
+            for attribute in attributes {
+                let AttributeKind::Unroll(unroll) = attribute else {
+                    continue;
+                };
+                // 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::Full => "llvm.loop.unroll.full",
+                        UnrollAttr::Never => "llvm.loop.unroll.disable",
+                        UnrollAttr::Count(_) => unreachable!(),
+                    };
+                    let unroll_meta = self.create_metadata(metadata_str.as_bytes());
+                    self.md_node_in_context(&[unroll_meta])
+                };
+                nodes.push(md_node);
+            }
+
+            if let [first, ..] = nodes[..] {
+                nodes.insert(0, first);
+
+                // 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

@@ -1091,9 +1091,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

@@ -984,6 +984,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,
@@ -1013,6 +1017,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::attrs::AttributeKind;
 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: &[AttributeKind],
     ) -> 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
         }
@@ -291,7 +293,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
@@ -359,7 +361,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
@@ -618,7 +620,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());
@@ -727,7 +729,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
@@ -861,7 +863,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);
@@ -935,7 +937,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();
@@ -1023,7 +1025,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
@@ -1133,7 +1135,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;
@@ -1577,7 +1579,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::attrs::AttributeKind;
 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: &[AttributeKind]) {
+        self.br(dest)
+    }
     fn cond_br(
         &mut self,
         cond: Self::Value,