Auto merge of rust-lang#127210 - jieyouxu:array-imm-noundef-param, r=<try>

Annotate eligible small immediate arguments with `noundef`

Retake of rust-lang#123425.

We try to annotate small (fits within target pointer width) aggregate arguments passed as immediates
(specifically casted as "appropriately sized integer type") with `noundef`.

Example:

```rs
#[no_mangle]
pub fn short_array_u64x1(v: [u64; 1]) -> [u64; 1] { v }
```

currently produces

```llvm
define i64 `@short_array_u64x1(i64` %0) ...
```

This PR changes that to

```llvm
define noundef i64 `@short_array_u64x1(i64` noundef %0) ...
```

The `noundef` attribute is added only if the immediate value has no padding. Specifically,
the conservative heuristic we use is to:

- Peel away layers of `#[repr(Rust)]` or `#[repr(transparent)]` wrappers if present
- Check for innermost simple arrays (whose element type is a primitive type) that can fit within
  target pointer width

Union immediates or otherwise anything that contains unions will not have `noundef` attribute
applied.

Closes rust-lang#123183.

cc `@/RalfJung` who pointed out various problems with the previous take, hopefully I addressed most of
them in this take.

r? `@ghost` (perf)

Author: bors

compiler/rustc_abi/src/lib.rs

@@ -101,6 +101,10 @@ pub struct ReprOptions {
 }
 
 impl ReprOptions {
+    pub fn rust(&self) -> bool {
+        !(self.simd() || self.c() || self.packed() || self.transparent() || self.linear())
+    }
+
     #[inline]
     pub fn simd(&self) -> bool {
         self.flags.contains(ReprFlags::IS_SIMD)

compiler/rustc_ty_utils/src/abi.rs

@@ -9,8 +9,8 @@ use rustc_middle::ty::{self, InstanceKind, Ty, TyCtxt};
 use rustc_session::config::OptLevel;
 use rustc_span::def_id::DefId;
 use rustc_target::abi::call::{
-    ArgAbi, ArgAttribute, ArgAttributes, ArgExtension, Conv, FnAbi, PassMode, Reg, RegKind,
-    RiscvInterruptKind,
+    ArgAbi, ArgAttribute, ArgAttributes, ArgExtension, CastTarget, Conv, FnAbi, PassMode, Reg,
+    RegKind, RiscvInterruptKind,
 };
 use rustc_target::abi::*;
 use rustc_target::spec::abi::Abi as SpecAbi;
@@ -787,6 +787,19 @@ fn fn_abi_adjust_for_abi<'tcx>(
                 // an LLVM aggregate type for this leads to bad optimizations,
                 // so we pick an appropriately sized integer type instead.
                 arg.cast_to(Reg { kind: RegKind::Integer, size });
+
+                // Now we see if we are allowed to annotate the small immediate argument with
+                // `noundef`. This is only legal for small aggregates that do not have padding. For
+                // instance, all unions are allowed `undef` so we must not annotate union (or
+                // anything that contain unions) immediates with `noundef`.
+                if can_annotate_small_immediate_argument_with_noundef(cx, arg.layout) {
+                    // Fixup arg attribute with `noundef`.
+                    let PassMode::Cast { ref mut cast, .. } = &mut arg.mode else {
+                        bug!("this cannot fail because of the previous cast_to `Reg`");
+                    };
+                    let box CastTarget { ref mut attrs, .. } = cast;
+                    attrs.set(ArgAttribute::NoUndef);
+                }
             }
 
             // If we deduced that this parameter was read-only, add that to the attribute list now.
@@ -822,6 +835,44 @@ fn fn_abi_adjust_for_abi<'tcx>(
     Ok(())
 }
 
+fn can_annotate_small_immediate_argument_with_noundef<'tcx>(
+    cx: &LayoutCx<'tcx, TyCtxt<'tcx>>,
+    outermost_layout: TyAndLayout<'tcx>,
+) -> bool {
+    fn allowed_repr_rust<'tcx>(candidate_ty: Ty<'tcx>) -> bool {
+        match candidate_ty.kind() {
+            ty::Adt(def, _) if def.repr().rust() && !def.is_union() => true,
+            ty::Array(elem_ty, _) if allowed_repr_rust(*elem_ty) => true,
+            t => t.is_primitive(),
+        }
+    }
+
+    fn is_transparent_or_rust_wrapper<'tcx>(layout: TyAndLayout<'tcx>) -> bool {
+        return layout.is_transparent::<LayoutCx<'tcx, TyCtxt<'tcx>>>()
+            || allowed_repr_rust(layout.ty);
+    }
+
+    if outermost_layout.ty.is_union() {
+        return false;
+    }
+
+    let mut innermost_layout = outermost_layout;
+    // Recursively peel away wrapper layers.
+    while is_transparent_or_rust_wrapper(innermost_layout) {
+        let Some((_, layout)) = innermost_layout.non_1zst_field(cx) else {
+            break;
+        };
+
+        if layout.ty.is_union() {
+            return false;
+        }
+
+        innermost_layout = layout;
+    }
+
+    allowed_repr_rust(innermost_layout.ty)
+}
+
 #[tracing::instrument(level = "debug", skip(cx))]
 fn make_thin_self_ptr<'tcx>(
     cx: &(impl HasTyCtxt<'tcx> + HasParamEnv<'tcx>),

tests/codegen/noundef-small-immediate-arguments.rs

@@ -0,0 +1,242 @@
+//! We would like to try add `noundef` to small immediate arguments (the heuristic here is if the
+//! immediate argument fits within target pointer width) where possible and legal. Adding `noundef`
+//! attribute is legal iff the immediate argument do not have padding (indeterminate value
+//! otherwise). Only simple immediates are considered currently (trivially no padding), but could be
+//! potentially expanded to other small aggregate immediates that do not have padding (subject to
+//! being able to correctly calculate "no padding").
+//!
+//! - We should recursively see through `#[repr(transparent)]` and `#[repr(Rust)]` layouts.
+//! - Unions cannot have `noundef` because all unions are currently allowed to be `undef`. This
+//!   property is "infectious", anything that contains unions also may not have `noundef` applied.
+
+// ignore-tidy-linelength
+
+#![crate_type = "lib"]
+
+//@ compile-flags: -C no-prepopulate-passes
+
+// We setup two revisions to check that `noundef` is only added when optimization is enabled.
+//@ revisions: NoOpt Opt
+//@ [NoOpt] compile-flags: -C opt-level=0
+//@ [Opt] compile-flags: -O
+
+// Presence of `noundef` depends on target pointer width (it's only applied when the immediate fits
+// within target pointer width).
+//@ only-64bit
+
+// -------------------------------------------------------------------------------------------------
+
+// # Positive test cases
+//
+// - Simple arrays of primitive types whose size fits within target pointer width (referred to as
+//   "simple arrays" for the following positive test cases).
+// - `#[repr(transparent)]` ADTs which eventually contain simple arrays.
+// - `#[repr(Rust)]` ADTs which eventually contain simple arrays. This relies on rustc layout
+//   behavior, and is not guaranteed by `#[repr(Rust)]`.
+
+// ## Simple arrays
+
+// NoOpt: define i64 @short_array_u64x1(i64 %{{.*}})
+// Opt: define noundef i64 @short_array_u64x1(i64 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u64x1(v: [u64; 1]) -> [u64; 1] {
+    v
+}
+
+// NoOpt: define i32 @short_array_u32x1(i32 %{{.*}})
+// Opt: define noundef i32 @short_array_u32x1(i32 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u32x1(v: [u32; 1]) -> [u32; 1] {
+    v
+}
+
+// NoOpt: define i64 @short_array_u32x2(i64 %{{.*}})
+// Opt: define noundef i64 @short_array_u32x2(i64 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u32x2(v: [u32; 2]) -> [u32; 2] {
+    v
+}
+
+// NoOpt: define i16 @short_array_u16x1(i16 %{{.*}})
+// Opt: define noundef i16 @short_array_u16x1(i16 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u16x1(v: [u16; 1]) -> [u16; 1] {
+    v
+}
+
+// NoOpt: define i32 @short_array_u16x2(i32 %{{.*}})
+// Opt: define noundef i32 @short_array_u16x2(i32 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u16x2(v: [u16; 2]) -> [u16; 2] {
+    v
+}
+
+// NoOpt: define i48 @short_array_u16x3(i48 %{{.*}})
+// Opt: define noundef i48 @short_array_u16x3(i48 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u16x3(v: [u16; 3]) -> [u16; 3] {
+    v
+}
+
+// NoOpt: define i64 @short_array_u16x4(i64 %{{.*}})
+// Opt: define noundef i64 @short_array_u16x4(i64 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u16x4(v: [u16; 4]) -> [u16; 4] {
+    v
+}
+
+// NoOpt: define i8 @short_array_u8x1(i8 %{{.*}})
+// Opt: define noundef i8 @short_array_u8x1(i8 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u8x1(v: [u8; 1]) -> [u8; 1] {
+    v
+}
+
+// NoOpt: define i16 @short_array_u8x2(i16 %{{.*}})
+// Opt: define noundef i16 @short_array_u8x2(i16 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u8x2(v: [u8; 2]) -> [u8; 2] {
+    v
+}
+
+// NoOpt: define i24 @short_array_u8x3(i24 %{{.*}})
+// Opt: define noundef i24 @short_array_u8x3(i24 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u8x3(v: [u8; 3]) -> [u8; 3] {
+    v
+}
+
+// NoOpt: define i64 @short_array_u8x8(i64 %{{.*}})
+// Opt: define noundef i64 @short_array_u8x8(i64 noundef %{{.*}})
+#[no_mangle]
+pub fn short_array_u8x8(v: [u8; 8]) -> [u8; 8] {
+    v
+}
+
+// ## Small `#[repr(transparent)]` wrappers
+
+#[repr(transparent)]
+pub struct TransparentWrapper([u8; 4]);
+
+// NoOpt: define i32 @repr_transparent_wrapper(i32 %{{.*}})
+// Opt: define noundef i32 @repr_transparent_wrapper(i32 noundef %{{.*}})
+#[no_mangle]
+pub fn repr_transparent_wrapper(v: TransparentWrapper) -> TransparentWrapper {
+    v
+}
+
+#[repr(transparent)]
+pub struct RecursiveTransparentWrapper(TransparentWrapper);
+
+// NoOpt: define i32 @recursive_repr_transparent_wrapper(i32 %{{.*}})
+// Opt: define noundef i32 @recursive_repr_transparent_wrapper(i32 noundef %{{.*}})
+#[no_mangle]
+pub fn recursive_repr_transparent_wrapper(
+    v: RecursiveTransparentWrapper,
+) -> RecursiveTransparentWrapper {
+    v
+}
+
+// ## Small `#[repr(Rust)]` wrappers
+//
+// Note that this relies on rustc self-consistency in handling simple `#[repr(Rust)]` wrappers, i.e.
+// that `struct Foo([u8; 4])` has the same layout as its sole inner member and that no additional
+// padding is introduced.
+
+pub struct ReprRustWrapper([u8; 4]);
+
+// NoOpt: define i32 @repr_rust_wrapper(i32 %{{.*}})
+// Opt: define noundef i32 @repr_rust_wrapper(i32 noundef %{{.*}})
+#[no_mangle]
+pub fn repr_rust_wrapper(v: ReprRustWrapper) -> ReprRustWrapper {
+    v
+}
+
+// ## Cases not handled
+//
+// - Aggregates that have no padding and fits within target pointer width which are not simple
+//   arrays. Potentially aggregates such as tuples `(u32, u32)`. This is left as an exercise to the
+//   reader (follow-up welcomed) :)
+
+// No `noundef` annotation on return `{i32, i32}`, but argument does (when optimizations enabled).
+// NoOpt: define { i32, i32 } @unhandled_small_pair_ret(i32 %v.0, i32 %v.1)
+// Opt: define { i32, i32 } @unhandled_small_pair_ret(i32 noundef %v.0, i32 noundef %v.1)
+#[no_mangle]
+pub fn unhandled_small_pair_ret(v: (u32, u32)) -> (u32, u32) {
+    v
+}
+
+// -------------------------------------------------------------------------------------------------
+
+// # Negative test cases ()
+//
+// - Other representations (not `transparent` or `Rust`)
+// - Unions cannot have `noundef` because they are allowed to be `undef`.
+// - Array of unions still contains unions, so they cannot have `noundef`
+// - Transparent unions are still unions, so they cannot have `noundef`
+
+// ## Other representations
+
+#[repr(C)]
+pub struct ReprCWrapper([u8; 4]);
+
+// NoOpt: define i32 @repr_c_immediate(i32 %0)
+// Opt: define i32 @repr_c_immediate(i32 %0)
+#[no_mangle]
+pub fn repr_c_immediate(v: ReprCWrapper) -> ReprCWrapper {
+    v
+}
+
+// ## Unions
+
+union U {
+    u1: u64,
+    u2: [u8; 4],
+}
+
+// All unions can be `undef`, must not have `noundef` as immediate argument.
+// NoOpt: define i64 @union_immediate(i64 %0)
+// Opt: define i64 @union_immediate(i64 %0)
+#[no_mangle]
+pub fn union_immediate(v: U) -> U {
+    v
+}
+
+// ## Array of unions
+//
+// Cannot have `noundef` because tainted by unions.
+
+union SmallButDangerous {
+    u1: [u8; 2],
+    u2: u16,
+}
+
+// NoOpt: define i16 @one_elem_array_of_unions(i16 %0)
+// Opt: define i16 @one_elem_array_of_unions(i16 %0)
+#[no_mangle]
+pub fn one_elem_array_of_unions(v: [SmallButDangerous; 1]) -> [SmallButDangerous; 1] {
+    v
+}
+
+// NoOpt: define i32 @two_elem_array_of_unions(i32 %0)
+// Opt: define i32 @two_elem_array_of_unions(i32 %0)
+#[no_mangle]
+pub fn two_elem_array_of_unions(v: [SmallButDangerous; 2]) -> [SmallButDangerous; 2] {
+    v
+}
+
+// # `#[repr(transparent)]` unions
+
+union Inner {
+    i1: u8,
+}
+
+#[repr(transparent)]
+pub struct TransparentUnionWrapper(Inner);
+
+// NoOpt: define i8 @repr_transparent_union(i8 %v)
+// Opt: define i8 @repr_transparent_union(i8 %v)
+#[no_mangle]
+pub fn repr_transparent_union(v: TransparentUnionWrapper) -> TransparentUnionWrapper {
+    v
+}