Auto merge of #111768 - oli-obk:pair_const_llvm, r=cjgillot

Optimize scalar and scalar pair representations loaded from ByRef in llvm

in https://github.com/rust-lang/rust/pull/105653 I noticed that we were generating suboptimal LLVM IR if we had a `ConstValue::ByRef` that could be represented by a `ScalarPair`. Before https://github.com/rust-lang/rust/pull/105653 this is probably rare, but after it, every slice will go down this suboptimal code path that requires LLVM to untangle a bunch of indirections and translate static allocations that are only used once to read a scalar pair from.

6 files changed, 120 insertions, 83 deletions

diff --git a/compiler/rustc_codegen_gcc/src/common.rs b/compiler/rustc_codegen_gcc/src/common.rs
index ac04b61a306..bad87db4732 100644
--- a/compiler/rustc_codegen_gcc/src/common.rs
+++ b/compiler/rustc_codegen_gcc/src/common.rs
@@ -1,17 +1,15 @@
 use gccjit::LValue;
 use gccjit::{RValue, Type, ToRValue};
-use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::{
     BaseTypeMethods,
     ConstMethods,
-    DerivedTypeMethods,
     MiscMethods,
     StaticMethods,
 };
 use rustc_middle::mir::Mutability;
-use rustc_middle::ty::layout::{TyAndLayout, LayoutOf};
+use rustc_middle::ty::layout::{LayoutOf};
 use rustc_middle::mir::interpret::{ConstAllocation, GlobalAlloc, Scalar};
-use rustc_target::abi::{self, HasDataLayout, Pointer, Size};
+use rustc_target::abi::{self, HasDataLayout, Pointer};
 
 use crate::consts::const_alloc_to_gcc;
 use crate::context::CodegenCx;
@@ -240,28 +238,26 @@ impl<'gcc, 'tcx> ConstMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         const_alloc_to_gcc(self, alloc)
     }
 
-    fn from_const_alloc(&self, layout: TyAndLayout<'tcx>, alloc: ConstAllocation<'tcx>, offset: Size) -> PlaceRef<'tcx, RValue<'gcc>> {
-        assert_eq!(alloc.inner().align, layout.align.abi);
-        let ty = self.type_ptr_to(layout.gcc_type(self));
-        let value =
-            if layout.size == Size::ZERO {
-                let value = self.const_usize(alloc.inner().align.bytes());
-                self.const_bitcast(value, ty)
-            }
-            else {
-                let init = const_alloc_to_gcc(self, alloc);
-                let base_addr = self.static_addr_of(init, alloc.inner().align, None);
-
-                let array = self.const_bitcast(base_addr, self.type_i8p());
-                let value = self.context.new_array_access(None, array, self.const_usize(offset.bytes())).get_address(None);
-                self.const_bitcast(value, ty)
-            };
-        PlaceRef::new_sized(value, layout)
-    }
-
     fn const_ptrcast(&self, val: RValue<'gcc>, ty: Type<'gcc>) -> RValue<'gcc> {
         self.context.new_cast(None, val, ty)
     }
+
+    fn const_bitcast(&self, value: RValue<'gcc>, typ: Type<'gcc>) -> RValue<'gcc> {
+        if value.get_type() == self.bool_type.make_pointer() {
+            if let Some(pointee) = typ.get_pointee() {
+                if pointee.dyncast_vector().is_some() {
+                    panic!()
+                }
+            }
+        }
+        // NOTE: since bitcast makes a value non-constant, don't bitcast if not necessary as some
+        // SIMD builtins require a constant value.
+        self.bitcast_if_needed(value, typ)
+    }
+    
+    fn const_ptr_byte_offset(&self, base_addr: Self::Value, offset: abi::Size) -> Self::Value {
+        self.context.new_array_access(None, base_addr, self.const_usize(offset.bytes())).get_address(None)
+    }
 }
 
 pub trait SignType<'gcc, 'tcx> {
diff --git a/compiler/rustc_codegen_gcc/src/consts.rs b/compiler/rustc_codegen_gcc/src/consts.rs
index 792ab8f890d..873f652e6f1 100644
--- a/compiler/rustc_codegen_gcc/src/consts.rs
+++ b/compiler/rustc_codegen_gcc/src/consts.rs
@@ -1,6 +1,6 @@
 #[cfg(feature = "master")]
 use gccjit::FnAttribute;
-use gccjit::{Function, GlobalKind, LValue, RValue, ToRValue, Type};
+use gccjit::{Function, GlobalKind, LValue, RValue, ToRValue};
 use rustc_codegen_ssa::traits::{BaseTypeMethods, ConstMethods, DerivedTypeMethods, StaticMethods};
 use rustc_middle::span_bug;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
@@ -16,21 +16,6 @@ use crate::context::CodegenCx;
 use crate::errors::InvalidMinimumAlignment;
 use crate::type_of::LayoutGccExt;
 
-impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
-    pub fn const_bitcast(&self, value: RValue<'gcc>, typ: Type<'gcc>) -> RValue<'gcc> {
-        if value.get_type() == self.bool_type.make_pointer() {
-            if let Some(pointee) = typ.get_pointee() {
-                if pointee.dyncast_vector().is_some() {
-                    panic!()
-                }
-            }
-        }
-        // NOTE: since bitcast makes a value non-constant, don't bitcast if not necessary as some
-        // SIMD builtins require a constant value.
-        self.bitcast_if_needed(value, typ)
-    }
-}
-
 fn set_global_alignment<'gcc, 'tcx>(cx: &CodegenCx<'gcc, 'tcx>, gv: LValue<'gcc>, mut align: Align) {
     // The target may require greater alignment for globals than the type does.
     // Note: GCC and Clang also allow `__attribute__((aligned))` on variables,
diff --git a/compiler/rustc_codegen_llvm/src/common.rs b/compiler/rustc_codegen_llvm/src/common.rs
index 9127fba388b..a3910fef954 100644
--- a/compiler/rustc_codegen_llvm/src/common.rs
+++ b/compiler/rustc_codegen_llvm/src/common.rs
@@ -8,16 +8,15 @@ use crate::type_of::LayoutLlvmExt;
 use crate::value::Value;
 
 use rustc_ast::Mutability;
-use rustc_codegen_ssa::mir::place::PlaceRef;
 use rustc_codegen_ssa::traits::*;
 use rustc_data_structures::stable_hasher::{Hash128, HashStable, StableHasher};
 use rustc_hir::def_id::DefId;
 use rustc_middle::bug;
 use rustc_middle::mir::interpret::{ConstAllocation, GlobalAlloc, Scalar};
-use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
+use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::cstore::{DllCallingConvention, DllImport, PeImportNameType};
-use rustc_target::abi::{self, AddressSpace, HasDataLayout, Pointer, Size};
+use rustc_target::abi::{self, AddressSpace, HasDataLayout, Pointer};
 use rustc_target::spec::Target;
 
 use libc::{c_char, c_uint};
@@ -307,38 +306,24 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         const_alloc_to_llvm(self, alloc)
     }
 
-    fn from_const_alloc(
-        &self,
-        layout: TyAndLayout<'tcx>,
-        alloc: ConstAllocation<'tcx>,
-        offset: Size,
-    ) -> PlaceRef<'tcx, &'ll Value> {
-        let alloc_align = alloc.inner().align;
-        assert_eq!(alloc_align, layout.align.abi);
-        let llty = self.type_ptr_to(layout.llvm_type(self));
-        let llval = if layout.size == Size::ZERO {
-            let llval = self.const_usize(alloc_align.bytes());
-            unsafe { llvm::LLVMConstIntToPtr(llval, llty) }
-        } else {
-            let init = const_alloc_to_llvm(self, alloc);
-            let base_addr = self.static_addr_of(init, alloc_align, None);
-
-            let llval = unsafe {
-                llvm::LLVMRustConstInBoundsGEP2(
-                    self.type_i8(),
-                    self.const_bitcast(base_addr, self.type_i8p()),
-                    &self.const_usize(offset.bytes()),
-                    1,
-                )
-            };
-            self.const_bitcast(llval, llty)
-        };
-        PlaceRef::new_sized(llval, layout)
-    }
-
     fn const_ptrcast(&self, val: &'ll Value, ty: &'ll Type) -> &'ll Value {
         consts::ptrcast(val, ty)
     }
+
+    fn const_bitcast(&self, val: &'ll Value, ty: &'ll Type) -> &'ll Value {
+        self.const_bitcast(val, ty)
+    }
+
+    fn const_ptr_byte_offset(&self, base_addr: Self::Value, offset: abi::Size) -> Self::Value {
+        unsafe {
+            llvm::LLVMRustConstInBoundsGEP2(
+                self.type_i8(),
+                self.const_bitcast(base_addr, self.type_i8p()),
+                &self.const_usize(offset.bytes()),
+                1,
+            )
+        }
+    }
 }
 
 /// Get the [LLVM type][Type] of a [`Value`].
diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs
index 2301c3ef13e..4000c9540ce 100644
--- a/compiler/rustc_codegen_ssa/src/mir/operand.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/operand.rs
@@ -8,10 +8,10 @@ use crate::traits::*;
 use crate::MemFlags;
 
 use rustc_middle::mir;
-use rustc_middle::mir::interpret::{ConstValue, Pointer, Scalar};
+use rustc_middle::mir::interpret::{alloc_range, ConstValue, Pointer, Scalar};
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::Ty;
-use rustc_target::abi::{Abi, Align, Size};
+use rustc_target::abi::{self, Abi, Align, Size};
 
 use std::fmt;
 
@@ -115,13 +115,82 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                 OperandValue::Pair(a_llval, b_llval)
             }
             ConstValue::ByRef { alloc, offset } => {
-                return bx.load_operand(bx.from_const_alloc(layout, alloc, offset));
+                return Self::from_const_alloc(bx, layout, alloc, offset);
             }
         };
 
         OperandRef { val, layout }
     }
 
+    fn from_const_alloc<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+        bx: &mut Bx,
+        layout: TyAndLayout<'tcx>,
+        alloc: rustc_middle::mir::interpret::ConstAllocation<'tcx>,
+        offset: Size,
+    ) -> Self {
+        let alloc_align = alloc.inner().align;
+        assert_eq!(alloc_align, layout.align.abi);
+        let ty = bx.type_ptr_to(bx.cx().backend_type(layout));
+
+        let read_scalar = |start, size, s: abi::Scalar, ty| {
+            let val = alloc
+                .0
+                .read_scalar(
+                    bx,
+                    alloc_range(start, size),
+                    /*read_provenance*/ matches!(s.primitive(), abi::Pointer(_)),
+                )
+                .unwrap();
+            bx.scalar_to_backend(val, s, ty)
+        };
+
+        // It may seem like all types with `Scalar` or `ScalarPair` ABI are fair game at this point.
+        // However, `MaybeUninit<u64>` is considered a `Scalar` as far as its layout is concerned --
+        // and yet cannot be represented by an interpreter `Scalar`, since we have to handle the
+        // case where some of the bytes are initialized and others are not. So, we need an extra
+        // check that walks over the type of `mplace` to make sure it is truly correct to treat this
+        // like a `Scalar` (or `ScalarPair`).
+        match layout.abi {
+            Abi::Scalar(s @ abi::Scalar::Initialized { .. }) => {
+                let size = s.size(bx);
+                assert_eq!(size, layout.size, "abi::Scalar size does not match layout size");
+                let val = read_scalar(Size::ZERO, size, s, ty);
+                OperandRef { val: OperandValue::Immediate(val), layout }
+            }
+            Abi::ScalarPair(
+                a @ abi::Scalar::Initialized { .. },
+                b @ abi::Scalar::Initialized { .. },
+            ) => {
+                let (a_size, b_size) = (a.size(bx), b.size(bx));
+                let b_offset = a_size.align_to(b.align(bx).abi);
+                assert!(b_offset.bytes() > 0);
+                let a_val = read_scalar(
+                    Size::ZERO,
+                    a_size,
+                    a,
+                    bx.scalar_pair_element_backend_type(layout, 0, true),
+                );
+                let b_val = read_scalar(
+                    b_offset,
+                    b_size,
+                    b,
+                    bx.scalar_pair_element_backend_type(layout, 1, true),
+                );
+                OperandRef { val: OperandValue::Pair(a_val, b_val), layout }
+            }
+            _ if layout.is_zst() => OperandRef::new_zst(bx, layout),
+            _ => {
+                // Neither a scalar nor scalar pair. Load from a place
+                let init = bx.const_data_from_alloc(alloc);
+                let base_addr = bx.static_addr_of(init, alloc_align, None);
+
+                let llval = bx.const_ptr_byte_offset(base_addr, offset);
+                let llval = bx.const_bitcast(llval, ty);
+                bx.load_operand(PlaceRef::new_sized(llval, layout))
+            }
+        }
+    }
+
     /// Asserts that this operand refers to a scalar and returns
     /// a reference to its value.
     pub fn immediate(self) -> V {
diff --git a/compiler/rustc_codegen_ssa/src/traits/consts.rs b/compiler/rustc_codegen_ssa/src/traits/consts.rs
index 61906302779..dc2fc396480 100644
--- a/compiler/rustc_codegen_ssa/src/traits/consts.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/consts.rs
@@ -1,8 +1,6 @@
 use super::BackendTypes;
-use crate::mir::place::PlaceRef;
 use rustc_middle::mir::interpret::{ConstAllocation, Scalar};
-use rustc_middle::ty::layout::TyAndLayout;
-use rustc_target::abi::{self, Size};
+use rustc_target::abi;
 
 pub trait ConstMethods<'tcx>: BackendTypes {
     // Constant constructors
@@ -30,12 +28,8 @@ pub trait ConstMethods<'tcx>: BackendTypes {
     fn const_data_from_alloc(&self, alloc: ConstAllocation<'tcx>) -> Self::Value;
 
     fn scalar_to_backend(&self, cv: Scalar, layout: abi::Scalar, llty: Self::Type) -> Self::Value;
-    fn from_const_alloc(
-        &self,
-        layout: TyAndLayout<'tcx>,
-        alloc: ConstAllocation<'tcx>,
-        offset: Size,
-    ) -> PlaceRef<'tcx, Self::Value>;
 
     fn const_ptrcast(&self, val: Self::Value, ty: Self::Type) -> Self::Value;
+    fn const_bitcast(&self, val: Self::Value, ty: Self::Type) -> Self::Value;
+    fn const_ptr_byte_offset(&self, val: Self::Value, offset: abi::Size) -> Self::Value;
 }
diff --git a/tests/codegen/const_scalar_pair.rs b/tests/codegen/const_scalar_pair.rs
new file mode 100644
index 00000000000..8f32c50b798
--- /dev/null
+++ b/tests/codegen/const_scalar_pair.rs
@@ -0,0 +1,8 @@
+// compile-flags: --crate-type=lib -Copt-level=0 -Zmir-opt-level=0 -C debuginfo=2
+
+#![feature(inline_const)]
+
+pub fn foo() -> (i32, i32) {
+    // CHECK: ret { i32, i32 } { i32 1, i32 2 }
+    const { (1, 2) }
+}