Auto merge of rust-lang#99574 - durin42:allocator-patch-redux, r=nikic

codegen: use new {re,de,}allocator annotations in llvm

This obviates the patch that teaches LLVM internals about
_rust_{re,de}alloc functions by putting annotations directly in the IR
for the optimizer.

The sole test change is required to anchor FileCheck to the body of the
`box_uninitialized` method, so it doesn't see the `allocalign` on
`__rust_alloc` and get mad about the string `alloca` showing up. Since I
was there anyway, I added some checks on the attributes to prove the
right attributes got set.

r? `@nikic`

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -13,7 +13,7 @@ use smallvec::SmallVec;
 
 use crate::attributes;
 use crate::llvm::AttributePlace::Function;
-use crate::llvm::{self, Attribute, AttributeKind, AttributePlace};
+use crate::llvm::{self, AllocKindFlags, Attribute, AttributeKind, AttributePlace};
 use crate::llvm_util;
 pub use rustc_attr::{InlineAttr, InstructionSetAttr, OptimizeAttr};
 
@@ -227,6 +227,10 @@ pub(crate) fn default_optimisation_attrs<'ll>(
     attrs
 }
 
+fn create_alloc_family_attr(llcx: &llvm::Context) -> &llvm::Attribute {
+    llvm::CreateAttrStringValue(llcx, "alloc-family", "__rust_alloc")
+}
+
 /// Composite function which sets LLVM attributes for function depending on its AST (`#[attribute]`)
 /// attributes.
 pub fn from_fn_attrs<'ll, 'tcx>(
@@ -309,11 +313,54 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         // Need this for AArch64.
         to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
     }
-    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR) {
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR)
+        || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR_ZEROED)
+    {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            // 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]);
+            to_add.push(llvm::CreateAllocSizeAttr(cx.llcx, 0));
+            let mut flags = AllocKindFlags::Alloc | AllocKindFlags::Aligned;
+            if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR) {
+                flags |= AllocKindFlags::Uninitialized;
+            } else {
+                flags |= AllocKindFlags::Zeroed;
+            }
+            to_add.push(llvm::CreateAllocKindAttr(cx.llcx, flags));
+        }
         // apply to return place instead of function (unlike all other attributes applied in this function)
         let no_alias = AttributeKind::NoAlias.create_attr(cx.llcx);
         attributes::apply_to_llfn(llfn, AttributePlace::ReturnValue, &[no_alias]);
     }
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::REALLOCATOR) {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            to_add.push(llvm::CreateAllocKindAttr(
+                cx.llcx,
+                AllocKindFlags::Realloc | AllocKindFlags::Aligned,
+            ));
+            // applies to argument place instead of function place
+            let allocated_pointer = AttributeKind::AllocatedPointer.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(0), &[allocated_pointer]);
+            // apply to argument place instead of function
+            let alloc_align = AttributeKind::AllocAlign.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(2), &[alloc_align]);
+            to_add.push(llvm::CreateAllocSizeAttr(cx.llcx, 3));
+        }
+        let no_alias = AttributeKind::NoAlias.create_attr(cx.llcx);
+        attributes::apply_to_llfn(llfn, AttributePlace::ReturnValue, &[no_alias]);
+    }
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::DEALLOCATOR) {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            to_add.push(llvm::CreateAllocKindAttr(cx.llcx, AllocKindFlags::Free));
+            // applies to argument place instead of function place
+            let allocated_pointer = AttributeKind::AllocatedPointer.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(0), &[allocated_pointer]);
+        }
+    }
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::CMSE_NONSECURE_ENTRY) {
         to_add.push(llvm::CreateAttrString(cx.llcx, "cmse_nonsecure_entry"));
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -193,6 +193,9 @@ pub enum AttributeKind {
     SanitizeMemTag = 34,
     NoCfCheck = 35,
     ShadowCallStack = 36,
+    AllocSize = 37,
+    AllocatedPointer = 38,
+    AllocAlign = 39,
 }
 
 /// LLVMIntPredicate
@@ -986,6 +989,22 @@ pub mod debuginfo {
     }
 }
 
+use bitflags::bitflags;
+// These values **must** match with LLVMRustAllocKindFlags
+bitflags! {
+    #[repr(transparent)]
+    #[derive(Default)]
+    pub struct AllocKindFlags : u64 {
+        const Unknown = 0;
+        const Alloc = 1;
+        const Realloc = 1 << 1;
+        const Free = 1 << 2;
+        const Uninitialized = 1 << 3;
+        const Zeroed = 1 << 4;
+        const Aligned = 1 << 5;
+    }
+}
+
 extern "C" {
     pub type ModuleBuffer;
 }
@@ -1193,6 +1212,8 @@ extern "C" {
     pub fn LLVMRustCreateByValAttr<'a>(C: &'a Context, ty: &'a Type) -> &'a Attribute;
     pub fn LLVMRustCreateStructRetAttr<'a>(C: &'a Context, ty: &'a Type) -> &'a Attribute;
     pub fn LLVMRustCreateUWTableAttr(C: &Context, async_: bool) -> &Attribute;
+    pub fn LLVMRustCreateAllocSizeAttr(C: &Context, size_arg: u32) -> &Attribute;
+    pub fn LLVMRustCreateAllocKindAttr(C: &Context, size_arg: u64) -> &Attribute;
 
     // Operations on functions
     pub fn LLVMRustGetOrInsertFunction<'a>(

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -95,6 +95,14 @@ pub fn CreateUWTableAttr(llcx: &Context, async_: bool) -> &Attribute {
     unsafe { LLVMRustCreateUWTableAttr(llcx, async_) }
 }
 
+pub fn CreateAllocSizeAttr(llcx: &Context, size_arg: u32) -> &Attribute {
+    unsafe { LLVMRustCreateAllocSizeAttr(llcx, size_arg) }
+}
+
+pub fn CreateAllocKindAttr(llcx: &Context, kind_arg: AllocKindFlags) -> &Attribute {
+    unsafe { LLVMRustCreateAllocKindAttr(llcx, kind_arg.bits()) }
+}
+
 #[derive(Copy, Clone)]
 pub enum AttributePlace {
     ReturnValue,

compiler/rustc_feature/src/builtin_attrs.rs

@@ -532,6 +532,9 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
 
     rustc_attr!(rustc_allocator, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
     rustc_attr!(rustc_allocator_nounwind, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
+    rustc_attr!(rustc_reallocator, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
+    rustc_attr!(rustc_deallocator, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
+    rustc_attr!(rustc_allocator_zeroed, Normal, template!(Word), WarnFollowing, IMPL_DETAIL),
     gated!(
         alloc_error_handler, Normal, template!(Word), WarnFollowing,
         experimental!(alloc_error_handler)

compiler/rustc_llvm/llvm-wrapper/LLVMWrapper.h

@@ -86,6 +86,9 @@ enum LLVMRustAttribute {
   SanitizeMemTag = 34,
   NoCfCheck = 35,
   ShadowCallStack = 36,
+  AllocSize = 37,
+  AllocatedPointer = 38,
+  AllocAlign = 39,
 };
 
 typedef struct OpaqueRustString *RustStringRef;

compiler/rustc_llvm/llvm-wrapper/RustWrapper.cpp

@@ -234,6 +234,14 @@ static Attribute::AttrKind fromRust(LLVMRustAttribute Kind) {
     return Attribute::SanitizeMemTag;
   case ShadowCallStack:
     return Attribute::ShadowCallStack;
+  case AllocSize:
+    return Attribute::AllocSize;
+#if LLVM_VERSION_GE(15, 0)
+  case AllocatedPointer:
+    return Attribute::AllocatedPointer;
+  case AllocAlign:
+    return Attribute::AllocAlign;
+#endif
   }
   report_fatal_error("bad AttributeKind");
 }
@@ -305,6 +313,67 @@ extern "C" LLVMAttributeRef LLVMRustCreateUWTableAttr(LLVMContextRef C, bool Asy
 #endif
 }
 
+extern "C" LLVMAttributeRef LLVMRustCreateAllocSizeAttr(LLVMContextRef C, uint32_t ElementSizeArg) {
+  return wrap(Attribute::getWithAllocSizeArgs(*unwrap(C), ElementSizeArg, None));
+}
+
+#if LLVM_VERSION_GE(15, 0)
+
+// These values **must** match ffi::AllocKindFlags.
+// It _happens_ to match the LLVM values of llvm::AllocFnKind,
+// but that's happenstance and we do explicit conversions before
+// passing them to LLVM.
+enum class LLVMRustAllocKindFlags : uint64_t {
+  Unknown = 0,
+  Alloc = 1,
+  Realloc = 1 << 1,
+  Free = 1 << 2,
+  Uninitialized = 1 << 3,
+  Zeroed = 1 << 4,
+  Aligned = 1 << 5,
+};
+
+static LLVMRustAllocKindFlags operator&(LLVMRustAllocKindFlags A, LLVMRustAllocKindFlags B) {
+  return static_cast<LLVMRustAllocKindFlags>(static_cast<uint64_t>(A) &
+                                      static_cast<uint64_t>(B));
+}
+
+static bool isSet(LLVMRustAllocKindFlags F) { return F != LLVMRustAllocKindFlags::Unknown; }
+
+static llvm::AllocFnKind allocKindFromRust(LLVMRustAllocKindFlags F) {
+  llvm::AllocFnKind AFK = llvm::AllocFnKind::Unknown;
+  if (isSet(F & LLVMRustAllocKindFlags::Alloc)) {
+    AFK |= llvm::AllocFnKind::Alloc;
+  }
+  if (isSet(F & LLVMRustAllocKindFlags::Realloc)) {
+    AFK |= llvm::AllocFnKind::Realloc;
+  }
+  if (isSet(F & LLVMRustAllocKindFlags::Free)) {
+    AFK |= llvm::AllocFnKind::Free;
+  }
+  if (isSet(F & LLVMRustAllocKindFlags::Uninitialized)) {
+    AFK |= llvm::AllocFnKind::Uninitialized;
+  }
+  if (isSet(F & LLVMRustAllocKindFlags::Zeroed)) {
+    AFK |= llvm::AllocFnKind::Zeroed;
+  }
+  if (isSet(F & LLVMRustAllocKindFlags::Aligned)) {
+    AFK |= llvm::AllocFnKind::Aligned;
+  }
+  return AFK;
+}
+#endif
+
+extern "C" LLVMAttributeRef LLVMRustCreateAllocKindAttr(LLVMContextRef C, uint64_t AllocKindArg) {
+#if LLVM_VERSION_GE(15, 0)
+  return wrap(Attribute::get(*unwrap(C), Attribute::AllocKind,
+      static_cast<uint64_t>(allocKindFromRust(static_cast<LLVMRustAllocKindFlags>(AllocKindArg)))));
+#else
+  report_fatal_error(
+      "allockind attributes are new in LLVM 15 and should not be used on older LLVMs");
+#endif
+}
+
 // Enable a fast-math flag
 //
 // https://llvm.org/docs/LangRef.html#fast-math-flags

compiler/rustc_middle/src/middle/codegen_fn_attrs.rs

@@ -50,7 +50,7 @@ bitflags! {
         /// the hot path.
         const COLD                      = 1 << 0;
         /// `#[rustc_allocator]`: a hint to LLVM that the pointer returned from this
-        /// function is never null.
+        /// function is never null and the function has no side effects other than allocating.
         const ALLOCATOR                 = 1 << 1;
         /// An indicator that function will never unwind. Will become obsolete
         /// once C-unwind is fully stabilized.
@@ -91,6 +91,12 @@ bitflags! {
         const NO_COVERAGE               = 1 << 15;
         /// `#[used(linker)]`: indicates that LLVM nor the linker can eliminate this function.
         const USED_LINKER               = 1 << 16;
+        /// `#[rustc_deallocator]`: a hint to LLVM that the function only deallocates memory.
+        const DEALLOCATOR               = 1 << 17;
+        /// `#[rustc_reallocator]`: a hint to LLVM that the function only reallocates memory.
+        const REALLOCATOR               = 1 << 18;
+        /// `#[rustc_allocator_zeroed]`: a hint to LLVM that the function only allocates zeroed memory.
+        const ALLOCATOR_ZEROED          = 1 << 19;
     }
 }
 

compiler/rustc_span/src/symbol.rs

@@ -1202,6 +1202,7 @@ symbols! {
         rustc,
         rustc_allocator,
         rustc_allocator_nounwind,
+        rustc_allocator_zeroed,
         rustc_allow_const_fn_unstable,
         rustc_allow_incoherent_impl,
         rustc_allowed_through_unstable_modules,
@@ -1214,6 +1215,7 @@ symbols! {
         rustc_const_stable,
         rustc_const_unstable,
         rustc_conversion_suggestion,
+        rustc_deallocator,
         rustc_def_path,
         rustc_diagnostic_item,
         rustc_diagnostic_macros,
@@ -1258,6 +1260,7 @@ symbols! {
         rustc_private,
         rustc_proc_macro_decls,
         rustc_promotable,
+        rustc_reallocator,
         rustc_regions,
         rustc_reservation_impl,
         rustc_serialize,

compiler/rustc_typeck/src/collect.rs

@@ -2775,6 +2775,12 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: DefId) -> CodegenFnAttrs {
             }
         } else if attr.has_name(sym::rustc_allocator_nounwind) {
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::NEVER_UNWIND;
+        } else if attr.has_name(sym::rustc_reallocator) {
+            codegen_fn_attrs.flags |= CodegenFnAttrFlags::REALLOCATOR;
+        } else if attr.has_name(sym::rustc_deallocator) {
+            codegen_fn_attrs.flags |= CodegenFnAttrFlags::DEALLOCATOR;
+        } else if attr.has_name(sym::rustc_allocator_zeroed) {
+            codegen_fn_attrs.flags |= CodegenFnAttrFlags::ALLOCATOR_ZEROED;
         } else if attr.has_name(sym::naked) {
             codegen_fn_attrs.flags |= CodegenFnAttrFlags::NAKED;
         } else if attr.has_name(sym::no_mangle) {

library/alloc/src/alloc.rs

@@ -25,15 +25,18 @@ extern "Rust" {
     // (the code expanding that attribute macro generates those functions), or to call
     // the default implementations in libstd (`__rdl_alloc` etc. in `library/std/src/alloc.rs`)
     // otherwise.
-    // The rustc fork of LLVM also special-cases these function names to be able to optimize them
+    // The rustc fork of LLVM 14 and earlier also special-cases these function names to be able to optimize them
     // like `malloc`, `realloc`, and `free`, respectively.
     #[rustc_allocator]
     #[rustc_allocator_nounwind]
     fn __rust_alloc(size: usize, align: usize) -> *mut u8;
+    #[cfg_attr(not(bootstrap), rustc_deallocator)]
     #[rustc_allocator_nounwind]
     fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize);
+    #[cfg_attr(not(bootstrap), rustc_reallocator)]
     #[rustc_allocator_nounwind]
     fn __rust_realloc(ptr: *mut u8, old_size: usize, align: usize, new_size: usize) -> *mut u8;
+    #[cfg_attr(not(bootstrap), rustc_allocator_zeroed)]
     #[rustc_allocator_nounwind]
     fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8;
 }

src/test/codegen/box-maybe-uninit-llvm14.rs

@@ -0,0 +1,26 @@
+// compile-flags: -O
+
+// Once we're done with llvm 14 and earlier, this test can be deleted.
+
+#![crate_type="lib"]
+
+use std::mem::MaybeUninit;
+
+// Boxing a `MaybeUninit` value should not copy junk from the stack
+#[no_mangle]
+pub fn box_uninitialized() -> Box<MaybeUninit<usize>> {
+    // CHECK-LABEL: @box_uninitialized
+    // CHECK-NOT: store
+    // CHECK-NOT: alloca
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: memset
+    Box::new(MaybeUninit::uninit())
+}
+
+// FIXME: add a test for a bigger box. Currently broken, see
+// https://github.com/rust-lang/rust/issues/58201.
+
+// Hide the LLVM 15+ `allocalign` attribute in the declaration of __rust_alloc
+// from the CHECK-NOT above. We don't check the attributes here because we can't rely
+// on all of them being set until LLVM 15.
+// CHECK: declare noalias{{.*}} @__rust_alloc(i{{[0-9]+}}, i{{[0-9]+.*}})

src/test/codegen/box-maybe-uninit.rs

@@ -1,4 +1,5 @@
 // compile-flags: -O
+// min-llvm-version: 15.0
 #![crate_type="lib"]
 
 use std::mem::MaybeUninit;
@@ -16,3 +17,9 @@ pub fn box_uninitialized() -> Box<MaybeUninit<usize>> {
 
 // FIXME: add a test for a bigger box. Currently broken, see
 // https://github.com/rust-lang/rust/issues/58201.
+
+// Hide the `allocalign` attribute in the declaration of __rust_alloc
+// from the CHECK-NOT above, and also verify the attributes got set reasonably.
+// CHECK: declare noalias ptr @__rust_alloc(i{{[0-9]+}}, i{{[0-9]+}} allocalign) unnamed_addr [[RUST_ALLOC_ATTRS:#[0-9]+]]
+
+// CHECK-DAG: attributes [[RUST_ALLOC_ATTRS]] = { {{.*}} allockind("alloc,uninitialized,aligned") allocsize(0) uwtable "alloc-family"="__rust_alloc" {{.*}} }