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