Pass `alloc-variant-zeroed` to LLVM

5 files changed, 36 insertions, 0 deletions

diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs
index c548f467583..984db43c67f 100644
--- a/compiler/rustc_codegen_llvm/src/attributes.rs
+++ b/compiler/rustc_codegen_llvm/src/attributes.rs
@@ -436,6 +436,16 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
         || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR_ZEROED)
     {
         to_add.push(create_alloc_family_attr(cx.llcx));
+        if let Some(zv) =
+            cx.tcx.get_attr(instance.def_id(), rustc_span::sym::rustc_allocator_zeroed_variant)
+            && let Some(name) = zv.value_str()
+        {
+            to_add.push(llvm::CreateAttrStringValue(
+                cx.llcx,
+                "alloc-variant-zeroed",
+                &mangle_internal_symbol(cx.tcx, name.as_str()),
+            ));
+        }
         // apply to argument place instead of function
         let alloc_align = AttributeKind::AllocAlign.create_attr(cx.llcx);
         attributes::apply_to_llfn(llfn, AttributePlace::Argument(1), &[alloc_align]);
diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
index ab6b8f92802..6cd0c43a3e0 100644
--- a/compiler/rustc_feature/src/builtin_attrs.rs
+++ b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -993,6 +993,10 @@ pub static BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_allocator_zeroed, Normal, template!(Word), WarnFollowing,
         EncodeCrossCrate::No,
     ),
+    rustc_attr!(
+        rustc_allocator_zeroed_variant, Normal, template!(NameValueStr: "function"), ErrorPreceding,
+        EncodeCrossCrate::Yes,
+    ),
     gated!(
         default_lib_allocator, Normal, template!(Word), WarnFollowing,
         EncodeCrossCrate::No, allocator_internals, experimental!(default_lib_allocator),
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index acbed7a9eed..21ba6c2779d 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -1816,6 +1816,7 @@ symbols! {
         rustc_align,
         rustc_allocator,
         rustc_allocator_zeroed,
+        rustc_allocator_zeroed_variant,
         rustc_allow_const_fn_unstable,
         rustc_allow_incoherent_impl,
         rustc_allowed_through_unstable_modules,
diff --git a/library/alloc/src/alloc.rs b/library/alloc/src/alloc.rs
index c9b98fa4e5a..76630a746dd 100644
--- a/library/alloc/src/alloc.rs
+++ b/library/alloc/src/alloc.rs
@@ -17,6 +17,7 @@ unsafe extern "Rust" {
     #[rustc_allocator]
     #[rustc_nounwind]
     #[rustc_std_internal_symbol]
+    #[rustc_allocator_zeroed_variant = "__rust_alloc_zeroed"]
     fn __rust_alloc(size: usize, align: usize) -> *mut u8;
     #[rustc_deallocator]
     #[rustc_nounwind]
diff --git a/tests/codegen-llvm/vec-calloc.rs b/tests/codegen-llvm/vec-calloc.rs
index d1c320ead01..15971bbfa00 100644
--- a/tests/codegen-llvm/vec-calloc.rs
+++ b/tests/codegen-llvm/vec-calloc.rs
@@ -1,4 +1,6 @@
+//@ revisions: normal llvm21
 //@ compile-flags: -Copt-level=3 -Z merge-functions=disabled
+//@ [llvm21] min-llvm-version: 21
 //@ only-x86_64
 
 #![crate_type = "lib"]
@@ -176,6 +178,24 @@ pub fn vec_option_i32(n: usize) -> Vec<Option<i32>> {
     vec![None; n]
 }
 
+// LLVM21-LABEL: @vec_array
+#[cfg(llvm21)]
+#[no_mangle]
+pub fn vec_array(n: usize) -> Vec<[u32; 1_000_000]> {
+    // LLVM21-NOT: call {{.*}}alloc::vec::from_elem
+    // LLVM21-NOT: call {{.*}}reserve
+    // LLVM21-NOT: call {{.*}}__rust_alloc(
+
+    // LLVM21: call {{.*}}__rust_alloc_zeroed(
+
+    // LLVM21-NOT: call {{.*}}alloc::vec::from_elem
+    // LLVM21-NOT: call {{.*}}reserve
+    // LLVM21-NOT: call {{.*}}__rust_alloc(
+
+    // LLVM21: ret void
+    vec![[0; 1_000_000]; 3]
+}
+
 // Ensure that __rust_alloc_zeroed gets the right attributes for LLVM to optimize it away.
 // CHECK: declare noalias noundef ptr @{{.*}}__rust_alloc_zeroed(i64 noundef, i64 allocalign noundef) unnamed_addr [[RUST_ALLOC_ZEROED_ATTRS:#[0-9]+]]