Skip to content

Commit

Permalink
Allow transmutes to produce OperandValues instead of always using…
Browse files Browse the repository at this point in the history
… `alloca`s

LLVM can usually optimize these away, but especially for things like transmutes of newtypes it's silly to generate the `alloc`+`store`+`load` at all when it's actually a nop at LLVM level.
  • Loading branch information
scottmcm committed Apr 5, 2023
1 parent 480068c commit 9aa9a84
Show file tree
Hide file tree
Showing 6 changed files with 381 additions and 74 deletions.
18 changes: 17 additions & 1 deletion compiler/rustc_codegen_ssa/src/mir/operand.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,26 @@ pub enum OperandValue<V> {
/// to be valid for the operand's lifetime.
/// The second value, if any, is the extra data (vtable or length)
/// which indicates that it refers to an unsized rvalue.
///
/// An `OperandValue` has this variant for types which are neither
/// `Immediate` nor `Pair`s. The backend value in this variant must be a
/// pointer to the *non*-immediate backend type. That pointee type is the
/// one returned by [`LayoutTypeMethods::backend_type`].
Ref(V, Option<V>, Align),
/// A single LLVM value.
/// A single LLVM immediate value.
///
/// An `OperandValue` *must* be this variant for any type for which
/// [`LayoutTypeMethods::is_backend_immediate`] returns `true`.
/// The backend value in this variant must be the *immediate* backend type,
/// as returned by [`LayoutTypeMethods::immediate_backend_type`].
Immediate(V),
/// A pair of immediate LLVM values. Used by fat pointers too.
///
/// An `OperandValue` *must* be this variant for any type for which
/// [`LayoutTypeMethods::is_backend_scalar_pair`] returns `true`.
/// The backend values in this variant must be the *immediate* backend types,
/// as returned by [`LayoutTypeMethods::scalar_pair_element_backend_type`]
/// with `immediate: true`.
Pair(V, V),
}

Expand Down
178 changes: 159 additions & 19 deletions compiler/rustc_codegen_ssa/src/mir/rvalue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use crate::MemFlags;
use rustc_middle::mir;
use rustc_middle::mir::Operand;
use rustc_middle::ty::cast::{CastTy, IntTy};
use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout};
use rustc_middle::ty::{self, adjustment::PointerCast, Instance, Ty, TyCtxt};
use rustc_span::source_map::{Span, DUMMY_SP};
use rustc_target::abi::{self, FIRST_VARIANT};
Expand Down Expand Up @@ -159,8 +159,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
debug_assert!(dst.layout.is_sized());

if src.layout.size != dst.layout.size
|| src.layout.abi == abi::Abi::Uninhabited
|| dst.layout.abi == abi::Abi::Uninhabited
|| src.layout.abi.is_uninhabited()
|| dst.layout.abi.is_uninhabited()
{
// In all of these cases it's UB to run this transmute, but that's
// known statically so might as well trap for it, rather than just
Expand All @@ -169,22 +169,20 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
return;
}

let size_in_bytes = src.layout.size.bytes();
if size_in_bytes == 0 {
// Nothing to write
if let Some(val) = self.codegen_transmute_operand(bx, src, dst.layout) {
val.store(bx, dst);
return;
}

match src.val {
OperandValue::Ref(src_llval, meta, src_align) => {
debug_assert_eq!(meta, None);
// For a place-to-place transmute, call `memcpy` directly so that
// both arguments get the best-available alignment information.
let bytes = bx.cx().const_usize(size_in_bytes);
let flags = MemFlags::empty();
bx.memcpy(dst.llval, dst.align, src_llval, src_align, bytes, flags);
OperandValue::Ref(..) => {
span_bug!(
self.mir.span,
"Operand path should have handled transmute \
from `Ref` {src:?} to place {dst:?}"
);
}
OperandValue::Immediate(_) | OperandValue::Pair(_, _) => {
OperandValue::Immediate(..) | OperandValue::Pair(..) => {
// When we have immediate(s), the alignment of the source is irrelevant,
// so we can store them using the destination's alignment.
let llty = bx.backend_type(src.layout);
Expand All @@ -194,6 +192,94 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
}
}

/// Attempts to transmute an `OperandValue` to another `OperandValue`.
///
/// Returns `None` for cases that can't work in that framework, such as for
/// `Immediate`->`Ref` that needs an `alloc` to get the location.
fn codegen_transmute_operand(
&mut self,
bx: &mut Bx,
operand: OperandRef<'tcx, Bx::Value>,
cast: TyAndLayout<'tcx>,
) -> Option<OperandValue<Bx::Value>> {
// Callers already checked that the layout sizes match
debug_assert_eq!(operand.layout.size, cast.size);

let operand_kind = self.value_kind(operand.layout);
let cast_kind = self.value_kind(cast);

match operand.val {
OperandValue::Ref(ptr, meta, align) => {
debug_assert_eq!(meta, None);
debug_assert!(matches!(operand_kind, OperandValueKind::Ref));
let cast_bty = bx.backend_type(cast);
let cast_ptr = bx.pointercast(ptr, bx.type_ptr_to(cast_bty));
let fake_place = PlaceRef::new_sized_aligned(cast_ptr, cast, align);
Some(bx.load_operand(fake_place).val)
}
OperandValue::Immediate(imm) => {
let OperandValueKind::Immediate(in_scalar) = operand_kind else {
bug!("Found {operand_kind:?} for operand {operand:?}");
};
if let OperandValueKind::Immediate(out_scalar) = cast_kind {
let cast_bty = bx.backend_type(cast);
Some(OperandValue::Immediate(Self::transmute_immediate(
bx, imm, in_scalar, out_scalar, cast_bty,
)))
} else {
None
}
}
OperandValue::Pair(imm_a, imm_b) => {
let OperandValueKind::Pair(in_a, in_b) = operand_kind else {
bug!("Found {operand_kind:?} for operand {operand:?}");
};
if let OperandValueKind::Pair(out_a, out_b) = cast_kind {
let out_a_ibty = bx.scalar_pair_element_backend_type(cast, 0, false);
let out_b_ibty = bx.scalar_pair_element_backend_type(cast, 1, false);
Some(OperandValue::Pair(
Self::transmute_immediate(bx, imm_a, in_a, out_a, out_a_ibty),
Self::transmute_immediate(bx, imm_b, in_b, out_b, out_b_ibty),
))
} else {
None
}
}
}
}

/// Transmutes one of the immediates from an [`OperandValue::Immediate`]
/// or an [`OperandValue::Pair`] to an immediate of the target type.
///
/// `to_backend_ty` must be the *non*-immediate backend type (so it will be
/// `i8`, not `i1`, for `bool`-like types.)
fn transmute_immediate(
bx: &mut Bx,
mut imm: Bx::Value,
from_scalar: abi::Scalar,
to_scalar: abi::Scalar,
to_backend_ty: Bx::Type,
) -> Bx::Value {
use abi::Primitive::*;
imm = bx.from_immediate(imm);
imm = match (from_scalar.primitive(), to_scalar.primitive()) {
(Int(..) | F32 | F64, Int(..) | F32 | F64) => bx.bitcast(imm, to_backend_ty),
(Pointer(..), Pointer(..)) => bx.pointercast(imm, to_backend_ty),
(Int(..), Pointer(..)) => bx.inttoptr(imm, to_backend_ty),
(Pointer(..), Int(..)) => bx.ptrtoint(imm, to_backend_ty),
(F32 | F64, Pointer(..)) => {
let int_imm = bx.bitcast(imm, bx.cx().type_isize());
bx.inttoptr(int_imm, to_backend_ty)
}
(Pointer(..), F32 | F64) => {
let int_imm = bx.ptrtoint(imm, bx.cx().type_isize());
bx.bitcast(int_imm, to_backend_ty)
}
};
imm = bx.to_immediate_scalar(imm, to_scalar);
imm
}

pub fn codegen_rvalue_unsized(
&mut self,
bx: &mut Bx,
Expand Down Expand Up @@ -396,7 +482,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
OperandValue::Immediate(newval)
}
mir::CastKind::Transmute => {
bug!("Transmute operand {:?} in `codegen_rvalue_operand`", operand);
self.codegen_transmute_operand(bx, operand, cast).unwrap_or_else(|| {
bug!("Unsupported transmute-as-operand of {operand:?} to {cast:?}");
})
}
};
OperandRef { val, layout: cast }
Expand Down Expand Up @@ -739,10 +827,36 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
pub fn rvalue_creates_operand(&self, rvalue: &mir::Rvalue<'tcx>, span: Span) -> bool {
match *rvalue {
mir::Rvalue::Cast(mir::CastKind::Transmute, ..) =>
// FIXME: Now that transmute is an Rvalue, it would be nice if
// it could create `Immediate`s for scalars, where possible.
false,
mir::Rvalue::Cast(mir::CastKind::Transmute, ref operand, cast_ty) => {
let operand_ty = operand.ty(self.mir, self.cx.tcx());
let cast_layout = self.cx.layout_of(self.monomorphize(cast_ty));
let operand_layout = self.cx.layout_of(self.monomorphize(operand_ty));
if operand_layout.size != cast_layout.size
|| operand_layout.abi.is_uninhabited()
|| cast_layout.abi.is_uninhabited()
{
// Send UB cases to the full form so the operand version can
// `bitcast` without worrying about malformed IR.
return false;
}

match (self.value_kind(operand_layout), self.value_kind(cast_layout)) {
// Can always load from a pointer as needed
(OperandValueKind::Ref, _) => true,

// Need to generate an `alloc` to get a pointer from an immediate
(OperandValueKind::Immediate(..) | OperandValueKind::Pair(..), OperandValueKind::Ref) => false,

// When we have scalar immediates, we can convert them as needed
(OperandValueKind::Immediate(..), OperandValueKind::Immediate(..)) |
(OperandValueKind::Pair(..), OperandValueKind::Pair(..)) => true,

// Send mixings between scalars and pairs through the memory route
// FIXME: Maybe this could use insertvalue/extractvalue instead?
(OperandValueKind::Immediate(..), OperandValueKind::Pair(..)) |
(OperandValueKind::Pair(..), OperandValueKind::Immediate(..)) => false,
}
}
mir::Rvalue::Ref(..) |
mir::Rvalue::CopyForDeref(..) |
mir::Rvalue::AddressOf(..) |
Expand All @@ -767,4 +881,30 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {

// (*) this is only true if the type is suitable
}

/// Gets which variant of [`OperandValue`] is expected for a particular type.
fn value_kind(&self, layout: TyAndLayout<'tcx>) -> OperandValueKind {
if self.cx.is_backend_immediate(layout) {
debug_assert!(!self.cx.is_backend_scalar_pair(layout));
OperandValueKind::Immediate(match layout.abi {
abi::Abi::Scalar(s) => s,
abi::Abi::Vector { element, .. } => element,
x => bug!("Couldn't translate {x:?} as backend immediate"),
})
} else if self.cx.is_backend_scalar_pair(layout) {
let abi::Abi::ScalarPair(s1, s2) = layout.abi else {
bug!("Couldn't translate {:?} as backend scalar pair", layout.abi)
};
OperandValueKind::Pair(s1, s2)
} else {
OperandValueKind::Ref
}
}
}

#[derive(Debug, Copy, Clone)]
enum OperandValueKind {
Ref,
Immediate(abi::Scalar),
Pair(abi::Scalar, abi::Scalar),
}
11 changes: 11 additions & 0 deletions compiler/rustc_codegen_ssa/src/traits/type_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -100,11 +100,22 @@ pub trait DerivedTypeMethods<'tcx>: BaseTypeMethods<'tcx> + MiscMethods<'tcx> {
impl<'tcx, T> DerivedTypeMethods<'tcx> for T where Self: BaseTypeMethods<'tcx> + MiscMethods<'tcx> {}

pub trait LayoutTypeMethods<'tcx>: Backend<'tcx> {
/// The backend type used for a rust type when it's in memory,
/// such as when it's stack-allocated or when it's being loaded or stored.
fn backend_type(&self, layout: TyAndLayout<'tcx>) -> Self::Type;
fn cast_backend_type(&self, ty: &CastTarget) -> Self::Type;
fn fn_decl_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Self::Type;
fn fn_ptr_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Self::Type;
fn reg_backend_type(&self, ty: &Reg) -> Self::Type;
/// The backend type used for a rust type when it's in an SSA register.
///
/// For nearly all types this is the same as the [`Self::backend_type`], however
/// `bool` (and other `0`-or-`1` values) are kept as [`BaseTypeMethods::type_i1`]
/// in registers but as [`BaseTypeMethods::type_i8`] in memory.
///
/// Converting values between the two different backend types is done using
/// [`from_immediate`](super::BuilderMethods::from_immediate) and
/// [`to_immediate_scalar`](super::BuilderMethods::to_immediate_scalar).
fn immediate_backend_type(&self, layout: TyAndLayout<'tcx>) -> Self::Type;
fn is_backend_immediate(&self, layout: TyAndLayout<'tcx>) -> bool;
fn is_backend_scalar_pair(&self, layout: TyAndLayout<'tcx>) -> bool;
Expand Down
Loading

0 comments on commit 9aa9a84

Please sign in to comment.