Changes to constant handling - faster deduplication, more compact representation

Author: FractalFir

src/common.rs

@@ -9,7 +9,6 @@ use rustc_middle::mir::Mutability;
 use rustc_middle::mir::interpret::{ConstAllocation, GlobalAlloc, Scalar};
 use rustc_middle::ty::layout::LayoutOf;
 
-use crate::consts::const_alloc_to_gcc;
 use crate::context::CodegenCx;
 use crate::type_of::LayoutGccExt;
 
@@ -46,12 +45,65 @@ impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
 }
 
 pub fn bytes_in_context<'gcc, 'tcx>(cx: &CodegenCx<'gcc, 'tcx>, bytes: &[u8]) -> RValue<'gcc> {
-    let context = &cx.context;
-    let byte_type = context.new_type::<u8>();
-    let typ = context.new_array_type(None, byte_type, bytes.len() as u64);
-    let elements: Vec<_> =
-        bytes.iter().map(|&byte| context.new_rvalue_from_int(byte_type, byte as i32)).collect();
-    context.new_array_constructor(None, typ, &elements)
+    // Instead of always using an array of bytes, use an array of larger integers of target endianness
+    // if possible. This reduces the amount of `rvalues` we use, which reduces memory usage significantly.
+    //
+    // FIXME(FractalFir): Consider using `global_set_initializer` instead. Before this is done, we need to confirm that
+    // `global_set_initializer` is more memory efficient than the current solution.
+    // `global_set_initializer` calls `global_set_initializer_rvalue` under the hood - does it generate an array of rvalues,
+    // or is it using a more efficient representation?
+    match bytes.len() % 8 {
+        0 => {
+            let context = &cx.context;
+            let byte_type = context.new_type::<u64>();
+            let typ = context.new_array_type(None, byte_type, bytes.len() as u64 / 8);
+            let elements: Vec<_> = bytes
+                .chunks_exact(8)
+                .map(|arr| {
+                    let arr: [u8; 8] = arr.try_into().unwrap();
+                    context.new_rvalue_from_long(
+                        byte_type,
+                        // Since we are representing arbitrary byte runs as integers, we need to follow the target
+                        // endianness.
+                        match cx.sess().target.options.endian {
+                            rustc_abi::Endian::Little => u64::from_le_bytes(arr) as i64,
+                            rustc_abi::Endian::Big => u64::from_be_bytes(arr) as i64,
+                        },
+                    )
+                })
+                .collect();
+            context.new_array_constructor(None, typ, &elements)
+        }
+        4 => {
+            let context = &cx.context;
+            let byte_type = context.new_type::<u32>();
+            let typ = context.new_array_type(None, byte_type, bytes.len() as u64 / 4);
+            let elements: Vec<_> = bytes
+                .chunks_exact(4)
+                .map(|arr| {
+                    let arr: [u8; 4] = arr.try_into().unwrap();
+                    context.new_rvalue_from_int(
+                        byte_type,
+                        match cx.sess().target.options.endian {
+                            rustc_abi::Endian::Little => u32::from_le_bytes(arr) as i32,
+                            rustc_abi::Endian::Big => u32::from_be_bytes(arr) as i32,
+                        },
+                    )
+                })
+                .collect();
+            context.new_array_constructor(None, typ, &elements)
+        }
+        _ => {
+            let context = cx.context;
+            let byte_type = context.new_type::<u8>();
+            let typ = context.new_array_type(None, byte_type, bytes.len() as u64);
+            let elements: Vec<_> = bytes
+                .iter()
+                .map(|&byte| context.new_rvalue_from_int(byte_type, byte as i32))
+                .collect();
+            context.new_array_constructor(None, typ, &elements)
+        }
+    }
 }
 
 pub fn type_is_pointer(typ: Type<'_>) -> bool {
@@ -212,7 +264,7 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
                 let alloc_id = prov.alloc_id();
                 let base_addr = match self.tcx.global_alloc(alloc_id) {
                     GlobalAlloc::Memory(alloc) => {
-                        let init = const_alloc_to_gcc(self, alloc);
+                        let init = self.const_data_from_alloc(alloc);
                         let alloc = alloc.inner();
                         let value = match alloc.mutability {
                             Mutability::Mut => self.static_addr_of_mut(init, alloc.align, None),
@@ -234,7 +286,7 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
                                 }),
                             )))
                             .unwrap_memory();
-                        let init = const_alloc_to_gcc(self, alloc);
+                        let init = self.const_data_from_alloc(alloc);
                         self.static_addr_of(init, alloc.inner().align, None)
                     }
                     GlobalAlloc::Static(def_id) => {
@@ -257,7 +309,19 @@ impl<'gcc, 'tcx> ConstCodegenMethods for CodegenCx<'gcc, 'tcx> {
     }
 
     fn const_data_from_alloc(&self, alloc: ConstAllocation<'_>) -> Self::Value {
-        const_alloc_to_gcc(self, alloc)
+        // We ignore the alignment for the purpose of deduping RValues
+        // The alignment is not handled / used in any way by `const_alloc_to_gcc`,
+        // so it is OK to overwrite it here.
+        let mut mock_alloc = alloc.inner().clone();
+        mock_alloc.align = rustc_abi::Align::MAX;
+        // Check if the rvalue is already in the cache - if so, just return it directly.
+        if let Some(res) = self.const_cache.borrow().get(&mock_alloc) {
+            return *res;
+        }
+        // Rvalue not in the cache - convert and add it.
+        let res = crate::consts::const_alloc_to_gcc_uncached(self, alloc);
+        self.const_cache.borrow_mut().insert(mock_alloc, res);
+        res
     }
 
     fn const_ptr_byte_offset(&self, base_addr: Self::Value, offset: abi::Size) -> Self::Value {

src/consts.rs

@@ -42,18 +42,14 @@ fn set_global_alignment<'gcc, 'tcx>(
 
 impl<'gcc, 'tcx> StaticCodegenMethods for CodegenCx<'gcc, 'tcx> {
     fn static_addr_of(&self, cv: RValue<'gcc>, align: Align, kind: Option<&str>) -> RValue<'gcc> {
-        // TODO(antoyo): implement a proper rvalue comparison in libgccjit instead of doing the
-        // following:
-        for (value, variable) in &*self.const_globals.borrow() {
-            if format!("{:?}", value) == format!("{:?}", cv) {
-                if let Some(global_variable) = self.global_lvalues.borrow().get(variable) {
-                    let alignment = align.bits() as i32;
-                    if alignment > global_variable.get_alignment() {
-                        global_variable.set_alignment(alignment);
-                    }
+        if let Some(variable) = self.const_globals.borrow().get(&cv) {
+            if let Some(global_variable) = self.global_lvalues.borrow().get(variable) {
+                let alignment = align.bits() as i32;
+                if alignment > global_variable.get_alignment() {
+                    global_variable.set_alignment(alignment);
                 }
-                return *variable;
             }
+            return *variable;
         }
         let global_value = self.static_addr_of_mut(cv, align, kind);
         #[cfg(feature = "master")]
@@ -299,8 +295,10 @@ impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
         global
     }
 }
-
-pub fn const_alloc_to_gcc<'gcc>(
+/// Converts a given const alloc to a gcc Rvalue, without any caching or deduplication.
+/// YOU SHOULD NOT call this function directly - that may break the semantics of Rust.
+/// Use `const_data_from_alloc` instead.
+pub(crate) fn const_alloc_to_gcc_uncached<'gcc>(
     cx: &CodegenCx<'gcc, '_>,
     alloc: ConstAllocation<'_>,
 ) -> RValue<'gcc> {
@@ -371,7 +369,7 @@ fn codegen_static_initializer<'gcc, 'tcx>(
     def_id: DefId,
 ) -> Result<(RValue<'gcc>, ConstAllocation<'tcx>), ErrorHandled> {
     let alloc = cx.tcx.eval_static_initializer(def_id)?;
-    Ok((const_alloc_to_gcc(cx, alloc), alloc))
+    Ok((cx.const_data_from_alloc(alloc), alloc))
 }
 
 fn check_and_apply_linkage<'gcc, 'tcx>(

src/context.rs

@@ -1,4 +1,5 @@
 use std::cell::{Cell, RefCell};
+use std::collections::HashMap;
 
 use gccjit::{
     Block, CType, Context, Function, FunctionPtrType, FunctionType, LValue, Location, RValue, Type,
@@ -9,6 +10,7 @@ use rustc_codegen_ssa::errors as ssa_errors;
 use rustc_codegen_ssa::traits::{BackendTypes, BaseTypeCodegenMethods, MiscCodegenMethods};
 use rustc_data_structures::base_n::{ALPHANUMERIC_ONLY, ToBaseN};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_middle::mir::interpret::Allocation;
 use rustc_middle::mir::mono::CodegenUnit;
 use rustc_middle::span_bug;
 use rustc_middle::ty::layout::{
@@ -30,6 +32,8 @@ use crate::common::SignType;
 
 #[cfg_attr(not(feature = "master"), allow(dead_code))]
 pub struct CodegenCx<'gcc, 'tcx> {
+    /// A cache of converted ConstAllocs
+    pub const_cache: RefCell<HashMap<Allocation, RValue<'gcc>>>,
     pub codegen_unit: &'tcx CodegenUnit<'tcx>,
     pub context: &'gcc Context<'gcc>,
 
@@ -222,6 +226,7 @@ impl<'gcc, 'tcx> CodegenCx<'gcc, 'tcx> {
         }
 
         let mut cx = Self {
+            const_cache: Default::default(),
             codegen_unit,
             context,
             current_func: RefCell::new(None),