Skip to content

Commit 858fea4

Browse files
matthiaskrgrmatthiaskrgr
authored
Rollup merge of #88868 - calebzulawski:feature/simd_bitmask, r=workingjubilee
Allow simd_bitmask to return byte arrays cc `@rust-lang/project-portable-simd` `@workingjubilee`
2 parents a8dff34 + fe8ae57 commit 858fea4

File tree

6 files changed

+147
-59
lines changed

6 files changed

+147
-59
lines changed

‎compiler/rustc_codegen_llvm/src/intrinsic.rs

Lines changed: 69 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, LayoutOf};
1919
use rustc_middle::ty::{self, Ty};
2020
use rustc_middle::{bug, span_bug};
2121
use rustc_span::{sym, symbol::kw, Span, Symbol};
22-
use rustc_target::abi::{self, HasDataLayout, Primitive};
22+
use rustc_target::abi::{self, Align, HasDataLayout, Primitive};
2323
use rustc_target::spec::{HasTargetSpec, PanicStrategy};
2424

2525
use std::cmp::Ordering;
@@ -857,28 +857,39 @@ fn generic_simd_intrinsic(
857857
let arg_tys = sig.inputs();
858858

859859
if name == sym::simd_select_bitmask {
860-
let in_ty = arg_tys[0];
861-
let m_len = match in_ty.kind() {
862-
// Note that this `.unwrap()` crashes for isize/usize, that's sort
863-
// of intentional as there's not currently a use case for that.
864-
ty::Int(i) => i.bit_width().unwrap(),
865-
ty::Uint(i) => i.bit_width().unwrap(),
866-
_ => return_error!("`{}` is not an integral type", in_ty),
867-
};
868860
require_simd!(arg_tys[1], "argument");
869-
let (v_len, _) = arg_tys[1].simd_size_and_type(bx.tcx());
870-
require!(
871-
// Allow masks for vectors with fewer than 8 elements to be
872-
// represented with a u8 or i8.
873-
m_len == v_len || (m_len == 8 && v_len < 8),
874-
"mismatched lengths: mask length `{}` != other vector length `{}`",
875-
m_len,
876-
v_len
877-
);
861+
let (len, _) = arg_tys[1].simd_size_and_type(bx.tcx());
862+
863+
let expected_int_bits = (len.max(8) - 1).next_power_of_two();
864+
let expected_bytes = len / 8 + ((len % 8 > 0) as u64);
865+
866+
let mask_ty = arg_tys[0];
867+
let mask = match mask_ty.kind() {
868+
ty::Int(i) if i.bit_width() == Some(expected_int_bits) => args[0].immediate(),
869+
ty::Uint(i) if i.bit_width() == Some(expected_int_bits) => args[0].immediate(),
870+
ty::Array(elem, len)
871+
if matches!(elem.kind(), ty::Uint(ty::UintTy::U8))
872+
&& len.try_eval_usize(bx.tcx, ty::ParamEnv::reveal_all())
873+
== Some(expected_bytes) =>
874+
{
875+
let place = PlaceRef::alloca(bx, args[0].layout);
876+
args[0].val.store(bx, place);
877+
let int_ty = bx.type_ix(expected_bytes * 8);
878+
let ptr = bx.pointercast(place.llval, bx.cx.type_ptr_to(int_ty));
879+
bx.load(int_ty, ptr, Align::ONE)
880+
}
881+
_ => return_error!(
882+
"invalid bitmask `{}`, expected `u{}` or `[u8; {}]`",
883+
mask_ty,
884+
expected_int_bits,
885+
expected_bytes
886+
),
887+
};
888+
878889
let i1 = bx.type_i1();
879-
let im = bx.type_ix(v_len);
880-
let i1xn = bx.type_vector(i1, v_len);
881-
let m_im = bx.trunc(args[0].immediate(), im);
890+
let im = bx.type_ix(len);
891+
let i1xn = bx.type_vector(i1, len);
892+
let m_im = bx.trunc(mask, im);
882893
let m_i1s = bx.bitcast(m_im, i1xn);
883894
return Ok(bx.select(m_i1s, args[1].immediate(), args[2].immediate()));
884895
}
@@ -1056,16 +1067,16 @@ fn generic_simd_intrinsic(
10561067

10571068
if name == sym::simd_bitmask {
10581069
// The `fn simd_bitmask(vector) -> unsigned integer` intrinsic takes a
1059-
// vector mask and returns an unsigned integer containing the most
1060-
// significant bit (MSB) of each lane.
1061-
1062-
// If the vector has less than 8 lanes, a u8 is returned with zeroed
1063-
// trailing bits.
1070+
// vector mask and returns the most significant bit (MSB) of each lane in the form
1071+
// of either:
1072+
// * an unsigned integer
1073+
// * an array of `u8`
1074+
// If the vector has less than 8 lanes, a u8 is returned with zeroed trailing bits.
1075+
//
1076+
// The bit order of the result depends on the byte endianness, LSB-first for little
1077+
// endian and MSB-first for big endian.
10641078
let expected_int_bits = in_len.max(8);
1065-
match ret_ty.kind() {
1066-
ty::Uint(i) if i.bit_width() == Some(expected_int_bits) => (),
1067-
_ => return_error!("bitmask `{}`, expected `u{}`", ret_ty, expected_int_bits),
1068-
}
1079+
let expected_bytes = expected_int_bits / 8 + ((expected_int_bits % 8 > 0) as u64);
10691080

10701081
// Integer vector <i{in_bitwidth} x in_len>:
10711082
let (i_xn, in_elem_bitwidth) = match in_elem.kind() {
@@ -1095,8 +1106,34 @@ fn generic_simd_intrinsic(
10951106
let i1xn = bx.trunc(i_xn_msb, bx.type_vector(bx.type_i1(), in_len));
10961107
// Bitcast <i1 x N> to iN:
10971108
let i_ = bx.bitcast(i1xn, bx.type_ix(in_len));
1098-
// Zero-extend iN to the bitmask type:
1099-
return Ok(bx.zext(i_, bx.type_ix(expected_int_bits)));
1109+
1110+
match ret_ty.kind() {
1111+
ty::Uint(i) if i.bit_width() == Some(expected_int_bits) => {
1112+
// Zero-extend iN to the bitmask type:
1113+
return Ok(bx.zext(i_, bx.type_ix(expected_int_bits)));
1114+
}
1115+
ty::Array(elem, len)
1116+
if matches!(elem.kind(), ty::Uint(ty::UintTy::U8))
1117+
&& len.try_eval_usize(bx.tcx, ty::ParamEnv::reveal_all())
1118+
== Some(expected_bytes) =>
1119+
{
1120+
// Zero-extend iN to the array lengh:
1121+
let ze = bx.zext(i_, bx.type_ix(expected_bytes * 8));
1122+
1123+
// Convert the integer to a byte array
1124+
let ptr = bx.alloca(bx.type_ix(expected_bytes * 8), Align::ONE);
1125+
bx.store(ze, ptr, Align::ONE);
1126+
let array_ty = bx.type_array(bx.type_i8(), expected_bytes);
1127+
let ptr = bx.pointercast(ptr, bx.cx.type_ptr_to(array_ty));
1128+
return Ok(bx.load(array_ty, ptr, Align::ONE));
1129+
}
1130+
_ => return_error!(
1131+
"cannot return `{}`, expected `u{}` or `[u8; {}]`",
1132+
ret_ty,
1133+
expected_int_bits,
1134+
expected_bytes
1135+
),
1136+
}
11001137
}
11011138

11021139
fn simd_simple_float_intrinsic(

‎src/test/ui/simd/intrinsic/generic-bitmask.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -51,19 +51,19 @@ fn main() {
5151
let _: u64 = simd_bitmask(m64);
5252

5353
let _: u16 = simd_bitmask(m2);
54-
//~^ ERROR bitmask `u16`, expected `u8`
54+
//~^ ERROR invalid monomorphization of `simd_bitmask` intrinsic
5555

5656
let _: u16 = simd_bitmask(m8);
57-
//~^ ERROR bitmask `u16`, expected `u8`
57+
//~^ ERROR invalid monomorphization of `simd_bitmask` intrinsic
5858

5959
let _: u32 = simd_bitmask(m16);
60-
//~^ ERROR bitmask `u32`, expected `u16`
60+
//~^ ERROR invalid monomorphization of `simd_bitmask` intrinsic
6161

6262
let _: u64 = simd_bitmask(m32);
63-
//~^ ERROR bitmask `u64`, expected `u32`
63+
//~^ ERROR invalid monomorphization of `simd_bitmask` intrinsic
6464

6565
let _: u128 = simd_bitmask(m64);
66-
//~^ ERROR bitmask `u128`, expected `u64`
66+
//~^ ERROR invalid monomorphization of `simd_bitmask` intrinsic
6767

6868
}
6969
}

‎src/test/ui/simd/intrinsic/generic-bitmask.stderr

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,28 @@
1-
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u16`, expected `u8`
1+
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: cannot return `u16`, expected `u8` or `[u8; 1]`
22
--> $DIR/generic-bitmask.rs:53:22
33
|
44
LL | let _: u16 = simd_bitmask(m2);
55
| ^^^^^^^^^^^^^^^^
66

7-
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u16`, expected `u8`
7+
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: cannot return `u16`, expected `u8` or `[u8; 1]`
88
--> $DIR/generic-bitmask.rs:56:22
99
|
1010
LL | let _: u16 = simd_bitmask(m8);
1111
| ^^^^^^^^^^^^^^^^
1212

13-
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u32`, expected `u16`
13+
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: cannot return `u32`, expected `u16` or `[u8; 2]`
1414
--> $DIR/generic-bitmask.rs:59:22
1515
|
1616
LL | let _: u32 = simd_bitmask(m16);
1717
| ^^^^^^^^^^^^^^^^^
1818

19-
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u64`, expected `u32`
19+
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: cannot return `u64`, expected `u32` or `[u8; 4]`
2020
--> $DIR/generic-bitmask.rs:62:22
2121
|
2222
LL | let _: u64 = simd_bitmask(m32);
2323
| ^^^^^^^^^^^^^^^^^
2424

25-
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u128`, expected `u64`
25+
error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: cannot return `u128`, expected `u64` or `[u8; 8]`
2626
--> $DIR/generic-bitmask.rs:65:23
2727
|
2828
LL | let _: u128 = simd_bitmask(m64);

‎src/test/ui/simd/intrinsic/generic-select.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@ struct b8x4(pub i8, pub i8, pub i8, pub i8);
2020

2121
#[repr(simd)]
2222
#[derive(Copy, Clone, PartialEq)]
23-
struct b8x8(pub i8, pub i8, pub i8, pub i8,
24-
pub i8, pub i8, pub i8, pub i8);
23+
struct b8x8(pub i8, pub i8, pub i8, pub i8, pub i8, pub i8, pub i8, pub i8);
2524

2625
extern "platform-intrinsic" {
2726
fn simd_select<T, U>(x: T, a: U, b: U) -> U;
@@ -50,15 +49,15 @@ fn main() {
5049
//~^ ERROR found non-SIMD `u32`
5150

5251
simd_select_bitmask(0u16, x, x);
53-
//~^ ERROR mask length `16` != other vector length `4`
54-
//
52+
//~^ ERROR invalid bitmask `u16`, expected `u8` or `[u8; 1]`
53+
5554
simd_select_bitmask(0u8, 1u32, 2u32);
5655
//~^ ERROR found non-SIMD `u32`
5756

5857
simd_select_bitmask(0.0f32, x, x);
59-
//~^ ERROR `f32` is not an integral type
58+
//~^ ERROR invalid bitmask `f32`, expected `u8` or `[u8; 1]`
6059

6160
simd_select_bitmask("x", x, x);
62-
//~^ ERROR `&str` is not an integral type
61+
//~^ ERROR invalid bitmask `&str`, expected `u8` or `[u8; 1]`
6362
}
6463
}

‎src/test/ui/simd/intrinsic/generic-select.stderr

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,47 @@
11
error[E0511]: invalid monomorphization of `simd_select` intrinsic: mismatched lengths: mask length `8` != other vector length `4`
2-
--> $DIR/generic-select.rs:40:9
2+
--> $DIR/generic-select.rs:39:9
33
|
44
LL | simd_select(m8, x, x);
55
| ^^^^^^^^^^^^^^^^^^^^^
66

77
error[E0511]: invalid monomorphization of `simd_select` intrinsic: mask element type is `u32`, expected `i_`
8-
--> $DIR/generic-select.rs:43:9
8+
--> $DIR/generic-select.rs:42:9
99
|
1010
LL | simd_select(x, x, x);
1111
| ^^^^^^^^^^^^^^^^^^^^
1212

1313
error[E0511]: invalid monomorphization of `simd_select` intrinsic: mask element type is `f32`, expected `i_`
14-
--> $DIR/generic-select.rs:46:9
14+
--> $DIR/generic-select.rs:45:9
1515
|
1616
LL | simd_select(z, z, z);
1717
| ^^^^^^^^^^^^^^^^^^^^
1818

1919
error[E0511]: invalid monomorphization of `simd_select` intrinsic: expected SIMD argument type, found non-SIMD `u32`
20-
--> $DIR/generic-select.rs:49:9
20+
--> $DIR/generic-select.rs:48:9
2121
|
2222
LL | simd_select(m4, 0u32, 1u32);
2323
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2424

25-
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: mismatched lengths: mask length `16` != other vector length `4`
26-
--> $DIR/generic-select.rs:52:9
25+
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: invalid bitmask `u16`, expected `u8` or `[u8; 1]`
26+
--> $DIR/generic-select.rs:51:9
2727
|
2828
LL | simd_select_bitmask(0u16, x, x);
2929
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
3030

3131
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: expected SIMD argument type, found non-SIMD `u32`
32-
--> $DIR/generic-select.rs:55:9
32+
--> $DIR/generic-select.rs:54:9
3333
|
3434
LL | simd_select_bitmask(0u8, 1u32, 2u32);
3535
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
3636

37-
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: `f32` is not an integral type
38-
--> $DIR/generic-select.rs:58:9
37+
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: invalid bitmask `f32`, expected `u8` or `[u8; 1]`
38+
--> $DIR/generic-select.rs:57:9
3939
|
4040
LL | simd_select_bitmask(0.0f32, x, x);
4141
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
4242

43-
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: `&str` is not an integral type
44-
--> $DIR/generic-select.rs:61:9
43+
error[E0511]: invalid monomorphization of `simd_select_bitmask` intrinsic: invalid bitmask `&str`, expected `u8` or `[u8; 1]`
44+
--> $DIR/generic-select.rs:60:9
4545
|
4646
LL | simd_select_bitmask("x", x, x);
4747
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

‎src/test/ui/simd/simd-bitmask.rs

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
//run-pass
2+
//ignore-endian-big behavior of simd_select_bitmask is endian-specific
3+
#![feature(repr_simd, platform_intrinsics)]
4+
5+
extern "platform-intrinsic" {
6+
fn simd_bitmask<T, U>(v: T) -> U;
7+
fn simd_select_bitmask<T, U>(m: T, a: U, b: U) -> U;
8+
}
9+
10+
#[derive(Copy, Clone)]
11+
#[repr(simd)]
12+
struct Simd<T, const N: usize>([T; N]);
13+
14+
fn main() {
15+
unsafe {
16+
let v = Simd::<i8, 4>([-1, 0, -1, 0]);
17+
let i: u8 = simd_bitmask(v);
18+
let a: [u8; 1] = simd_bitmask(v);
19+
20+
assert_eq!(i, 0b0101);
21+
assert_eq!(a, [0b0101]);
22+
23+
let v = Simd::<i8, 16>([0, 0, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0, -1, 0]);
24+
let i: u16 = simd_bitmask(v);
25+
let a: [u8; 2] = simd_bitmask(v);
26+
27+
assert_eq!(i, 0b0101000000001100);
28+
assert_eq!(a, [0b1100, 0b01010000]);
29+
}
30+
31+
unsafe {
32+
let a = Simd::<i32, 8>([0, 1, 2, 3, 4, 5, 6, 7]);
33+
let b = Simd::<i32, 8>([8, 9, 10, 11, 12, 13, 14, 15]);
34+
let e = [0, 9, 2, 11, 12, 13, 14, 15];
35+
36+
let r = simd_select_bitmask(0b0101u8, a, b);
37+
assert_eq!(r.0, e);
38+
39+
let r = simd_select_bitmask([0b0101u8], a, b);
40+
assert_eq!(r.0, e);
41+
42+
let a = Simd::<i32, 16>([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]);
43+
let b = Simd::<i32, 16>([16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]);
44+
let e = [16, 17, 2, 3, 20, 21, 22, 23, 24, 25, 26, 27, 12, 29, 14, 31];
45+
46+
let r = simd_select_bitmask(0b0101000000001100u16, a, b);
47+
assert_eq!(r.0, e);
48+
49+
let r = simd_select_bitmask([0b1100u8, 0b01010000u8], a, b);
50+
assert_eq!(r.0, e);
51+
}
52+
}

0 commit comments

Comments
 (0)