Pass alloc-variant-zeroed to LLVM

Author: clubby789

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",
+                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]);

compiler/rustc_feature/src/builtin_attrs.rs

@@ -774,6 +774,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),

compiler/rustc_span/src/symbol.rs

@@ -1812,6 +1812,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,

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]

tests/codegen-llvm/vec-calloc.rs

@@ -176,6 +176,23 @@ pub fn vec_option_i32(n: usize) -> Vec<Option<i32>> {
     vec![None; n]
 }
 
+// CHECK-LABEL: @vec_array
+#[no_mangle]
+pub fn vec_array(n: usize) -> Vec<[u32; 1_000_000]> {
+    // CHECK-NOT: call {{.*}}alloc::vec::from_elem
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: call {{.*}}__rust_alloc(
+
+    // CHECK: call {{.*}}__rust_alloc_zeroed(
+
+    // CHECK-NOT: call {{.*}}alloc::vec::from_elem
+    // CHECK-NOT: call {{.*}}reserve
+    // CHECK-NOT: call {{.*}}__rust_alloc(
+
+    // CHECK: 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]+]]