codegen: add noundef for small immediate arguments with no padding

Add `noundef` metadata for arguments that:

1. are passed as immediates
2. fit in target pointer width
3. has no padding

Author: jieyouxu

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>),