Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make repr(packed) vectors work with SIMD intrinsics #125311

Merged
merged 3 commits into from
Jun 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 53 additions & 1 deletion compiler/rustc_codegen_llvm/src/intrinsic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -480,8 +480,60 @@ impl<'ll, 'tcx> IntrinsicCallMethods<'tcx> for Builder<'_, 'll, 'tcx> {
}

_ if name.as_str().starts_with("simd_") => {
// Unpack non-power-of-2 #[repr(packed, simd)] arguments.
// This gives them the expected layout of a regular #[repr(simd)] vector.
let mut loaded_args = Vec::new();
for (ty, arg) in arg_tys.iter().zip(args) {
loaded_args.push(
// #[repr(packed, simd)] vectors are passed like arrays (as references,
// with reduced alignment and no padding) rather than as immediates.
// We can use a vector load to fix the layout and turn the argument
// into an immediate.
if ty.is_simd()
&& let OperandValue::Ref(place) = arg.val
{
let (size, elem_ty) = ty.simd_size_and_type(self.tcx());
let elem_ll_ty = match elem_ty.kind() {
ty::Float(f) => self.type_float_from_ty(*f),
ty::Int(i) => self.type_int_from_ty(*i),
ty::Uint(u) => self.type_uint_from_ty(*u),
ty::RawPtr(_, _) => self.type_ptr(),
_ => unreachable!(),
};
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a function for this, Ty -> Type?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would have to be defined in cg_llvm instead of on Ty or TyKind, since each codegen backend would have to accept them and convert it into its local "type".

Copy link
Member Author

@calebzulawski calebzulawski May 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I was wondering if this exists somewhere in the builder traits so it could be defined per backend.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm. I looked around and I think it should only live in this file, because that's a Very Dicey thing to do in the general case, but perfectly reasonable here, and all of the uses of the from_ty family are here anyways.

Copy link
Member

@workingjubilee workingjubilee May 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

( It's possible I'm wrong about the diceyness... it seems to me that in most other cases you'd want to equate types carefully and be mindful... but I'm still ambiently unsure about whether this is the correct level of abstraction. )

let loaded =
self.load_from_place(self.type_vector(elem_ll_ty, size), place);
Comment on lines +503 to +504
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My other question I guess is why we need to do this instead of e.g. just letting the fact the type is Copy in all the cases we care about take over.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it really matters that it's Copy, the issue is that LLVM is expecting <n x ty> but is instead getting a (possibly underaligned) pointer

OperandRef::from_immediate_or_packed_pair(self, loaded, arg.layout)
} else {
*arg
},
);
}

let llret_ty = if ret_ty.is_simd()
&& let abi::Abi::Aggregate { .. } = self.layout_of(ret_ty).layout.abi
{
let (size, elem_ty) = ret_ty.simd_size_and_type(self.tcx());
let elem_ll_ty = match elem_ty.kind() {
ty::Float(f) => self.type_float_from_ty(*f),
ty::Int(i) => self.type_int_from_ty(*i),
ty::Uint(u) => self.type_uint_from_ty(*u),
ty::RawPtr(_, _) => self.type_ptr(),
_ => unreachable!(),
};
self.type_vector(elem_ll_ty, size)
} else {
llret_ty
};

match generic_simd_intrinsic(
self, name, callee_ty, fn_args, args, ret_ty, llret_ty, span,
self,
name,
callee_ty,
fn_args,
&loaded_args,
ret_ty,
llret_ty,
span,
) {
Ok(llval) => llval,
Err(()) => return Ok(()),
Expand Down
44 changes: 44 additions & 0 deletions tests/codegen/simd/packed-simd-alignment.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
//@ compile-flags: -Cno-prepopulate-passes

#![crate_type = "lib"]
#![feature(repr_simd, core_intrinsics)]
// make sure that codegen emits correctly-aligned loads and stores for repr(packed, simd) types
// the alignment of a load should be no less than T, and no more than the size of the vector type
use std::intrinsics::simd as intrinsics;

#[derive(Copy, Clone)]
#[repr(packed, simd)]
struct f32x3([f32; 3]);

#[derive(Copy, Clone)]
#[repr(packed, simd)]
struct f32x4([f32; 4]);

// CHECK-LABEL: load_f32x3
#[no_mangle]
pub fn load_f32x3(floats: &f32x3) -> f32x3 {
// FIXME: Is a memcpy really the best we can do?
// CHECK: @llvm.memcpy.{{.*}}ptr align 4 {{.*}}ptr align 4
*floats
}

// CHECK-LABEL: load_f32x4
#[no_mangle]
pub fn load_f32x4(floats: &f32x4) -> f32x4 {
// CHECK: load <4 x float>, ptr %{{[a-z0-9_]*}}, align {{4|8|16}}
*floats
}

// CHECK-LABEL: add_f32x3
#[no_mangle]
pub fn add_f32x3(x: f32x3, y: f32x3) -> f32x3 {
// CHECK: load <3 x float>, ptr %{{[a-z0-9_]*}}, align 4
unsafe { intrinsics::simd_add(x, y) }
}

// CHECK-LABEL: add_f32x4
#[no_mangle]
pub fn add_f32x4(x: f32x4, y: f32x4) -> f32x4 {
// CHECK: load <4 x float>, ptr %{{[a-z0-9_]*}}, align {{4|8|16}}
unsafe { intrinsics::simd_add(x, y) }
}
21 changes: 6 additions & 15 deletions tests/ui/simd/repr_packed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,6 @@
#[repr(simd, packed)]
struct Simd<T, const N: usize>([T; N]);

#[repr(simd)]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe a silly question, but is the plan ever to use non-packed simd? Should repr(simd) just mean what this PR does, now?

After all, the existing things like __m128 are all power-of-two length, so would still get the alignment they do today. And I, at least, find it really confusing that "packed" [u32; 8] actually has 32-byte alignment.

(repr(simd) isn't on stabilization track, and if someone does want a more-aligned non-PoT simd vector they can put it into an aligned newtype.)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, I think we should experiment with that separately, after this PR, so that we can still back out of this path in case we find this actually hits a shitton of LLVM errors in codegen on platforms that aren't x86-64.

I agree having a single handling would have nice qualities.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, "sure, but in a future PR" seems like a reasonable answer 👍

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there is a non-zero possibility that we end up with both types--one packed and one not. Notably, adding repr(packed) changes the (internal rustc) ABI from "vector" to "aggregate" and involves an extra load. I could see a situation where someone doesn't mind the padding and wants the (chance of) slightly better codegen opportunity.

It's of course slightly odd that it's not byte-aligned, but repr(packed(N)) does take an alignment argument, with repr(packed, simd) it's just a different default alignment.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

right. as far as I am concerned what we're really doing here is altering a very-poorly-defined lang item, since repr(simd) and repr(simd, packed) are Basically Magic that people can't really use directly, and what we want to think about is how to describe something that we can then expose to people.

struct FullSimd<T, const N: usize>([T; N]);

fn check_size_align<T, const N: usize>() {
use std::mem;
assert_eq!(mem::size_of::<Simd<T, N>>(), mem::size_of::<[T; N]>());
Expand Down Expand Up @@ -39,21 +36,15 @@ fn main() {
check_ty::<f64>();

unsafe {
// powers-of-two have no padding and work as usual
// powers-of-two have no padding and have the same layout as #[repr(simd)]
let x: Simd<f64, 4> =
simd_add(Simd::<f64, 4>([0., 1., 2., 3.]), Simd::<f64, 4>([2., 2., 2., 2.]));
assert_eq!(std::mem::transmute::<_, [f64; 4]>(x), [2., 3., 4., 5.]);

// non-powers-of-two have padding and need to be expanded to full vectors
fn load<T, const N: usize>(v: Simd<T, N>) -> FullSimd<T, N> {
unsafe {
let mut tmp = core::mem::MaybeUninit::<FullSimd<T, N>>::uninit();
std::ptr::copy_nonoverlapping(&v as *const _, tmp.as_mut_ptr().cast(), 1);
tmp.assume_init()
}
}
let x: FullSimd<f64, 3> =
simd_add(load(Simd::<f64, 3>([0., 1., 2.])), load(Simd::<f64, 3>([2., 2., 2.])));
assert_eq!(x.0, [2., 3., 4.]);
// non-powers-of-two should have padding (which is removed by #[repr(packed)]),
// but the intrinsic handles it
let x: Simd<f64, 3> = simd_add(Simd::<f64, 3>([0., 1., 2.]), Simd::<f64, 3>([2., 2., 2.]));
let arr: [f64; 3] = x.0;
assert_eq!(arr, [2., 3., 4.]);
}
}
Loading