Add unchecked methods, fix _mm_extract_epi* return types (rust-lang#223)

* Adds extract_unchecked + replace_unchecked + len (rust-lang#222 )

* [x86] Fixes the return types + uses extract_unchecked for:
  * _mm_extract_epi8
  * _mm_extract_epi16
  * _mm256_extract_epi8
  * _mm256_extract_epi16

* Minor changes to the other extract_epi* intrinsics for style consistency
These should now zero-extend the extracted int and behave appropriately. An old typo makes these a bit confusing, See this llvm issue.

Author: Tony Sifkarovski

coresimd/src/macros.rs

@@ -30,6 +30,11 @@ macro_rules! define_impl {
                 $name($($elname),*)
             }
 
+            #[inline(always)]
+            pub fn len() -> i32 {
+                $nelems
+            }
+
             #[inline(always)]
             pub fn splat(value: $elemty) -> $name {
                 $name($({
@@ -42,13 +47,27 @@ macro_rules! define_impl {
             #[inline(always)]
             pub fn extract(self, idx: u32) -> $elemty {
                 assert!(idx < $nelems);
-                unsafe { simd_extract(self, idx) }
+                unsafe { self.extract_unchecked(idx) }
+            }
+
+            #[inline(always)]
+            pub unsafe fn extract_unchecked(self, idx: u32) -> $elemty {
+                simd_extract(self, idx)
             }
 
             #[inline(always)]
             pub fn replace(self, idx: u32, val: $elemty) -> $name {
                 assert!(idx < $nelems);
-                unsafe { simd_insert(self, idx, val) }
+                unsafe { self.replace_unchecked(idx, val) }
+            }
+
+            #[inline(always)]
+            pub unsafe fn replace_unchecked(
+                self,
+                idx: u32,
+                val: $elemty,
+            ) -> $name {
+                simd_insert(self, idx, val)
             }
 
             #[inline(always)]

coresimd/src/x86/i586/avx.rs

@@ -863,36 +863,44 @@ pub unsafe fn _mm256_extractf128_si256(a: __m256i, imm8: i32) -> __m128i {
     __m128i::from(dst)
 }
 
-/// Extract an 8-bit integer from `a`, selected with `imm8`.
+/// Extract an 8-bit integer from `a`, selected with `imm8`. Returns a 32-bit
+/// integer containing the zero-extended integer data.
+/// See: https://reviews.llvm.org/D20468
 #[inline(always)]
 #[target_feature = "+avx"]
 // This intrinsic has no corresponding instruction.
 pub unsafe fn _mm256_extract_epi8(a: i8x32, imm8: i32) -> i32 {
-    a.extract(imm8 as u32 & 31) as i32
+    let imm8 = (imm8 & 31) as u32;
+    (a.extract_unchecked(imm8) as i32) & 0xFF
 }
 
-/// Extract a 16-bit integer from `a`, selected with `imm8`.
+/// Extract a 16-bit integer from `a`, selected with `imm8`. Returns a 32-bit
+/// integer containing the zero-extended integer data.
+/// See: https://reviews.llvm.org/D20468
 #[inline(always)]
 #[target_feature = "+avx"]
 // This intrinsic has no corresponding instruction.
 pub unsafe fn _mm256_extract_epi16(a: i16x16, imm8: i32) -> i32 {
-    a.extract(imm8 as u32 & 15) as i32
+    let imm8 = (imm8 & 15) as u32;
+    (a.extract_unchecked(imm8) as i32) & 0xFFFF
 }
 
 /// Extract a 32-bit integer from `a`, selected with `imm8`.
 #[inline(always)]
 #[target_feature = "+avx"]
 // This intrinsic has no corresponding instruction.
 pub unsafe fn _mm256_extract_epi32(a: i32x8, imm8: i32) -> i32 {
-    a.extract(imm8 as u32 & 7) as i32
+    let imm8 = (imm8 & 7) as u32;
+    a.extract_unchecked(imm8)
 }
 
 /// Extract a 64-bit integer from `a`, selected with `imm8`.
 #[inline(always)]
 #[target_feature = "+avx"]
 // This intrinsic has no corresponding instruction.
-pub unsafe fn _mm256_extract_epi64(a: i64x4, imm8: i32) -> i32 {
-    a.extract(imm8 as u32 & 3) as i32
+pub unsafe fn _mm256_extract_epi64(a: i64x4, imm8: i32) -> i64 {
+    let imm8 = (imm8 & 3) as u32;
+    a.extract_unchecked(imm8)
 }
 
 /// Zero the contents of all XMM or YMM registers.
@@ -3142,28 +3150,34 @@ mod tests {
     unsafe fn _mm256_extract_epi8() {
         #[cfg_attr(rustfmt, rustfmt_skip)]
         let a = i8x32::new(
-            1, 2, 3, 4, 5, 6, 7, 8,
-            9, 10, 11, 12, 13, 14, 15, 16,
-            17, 18, 19, 20, 21, 22, 23, 24,
-            25, 26, 27, 28, 29, 30, 31, 32,
+            -1, 1, 2, 3, 4, 5, 6, 7,
+            8, 9, 10, 11, 12, 13, 14, 15,
+            16, 17, 18, 19, 20, 21, 22, 23,
+            24, 25, 26, 27, 28, 29, 30, 31
         );
-        let r = avx::_mm256_extract_epi8(a, 0);
-        assert_eq!(r, 1);
+        let r1 = avx::_mm256_extract_epi8(a, 0);
+        let r2 = avx::_mm256_extract_epi8(a, 35);
+        assert_eq!(r1, 0xFF);
+        assert_eq!(r2, 3);
     }
 
     #[simd_test = "avx"]
     unsafe fn _mm256_extract_epi16() {
         let a =
-            i16x16::new(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
-        let r = avx::_mm256_extract_epi16(a, 0);
-        assert_eq!(r, 0);
+            i16x16::new(-1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+        let r1 = avx::_mm256_extract_epi16(a, 0);
+        let r2 = avx::_mm256_extract_epi16(a, 19);
+        assert_eq!(r1, 0xFFFF);
+        assert_eq!(r2, 3);
     }
 
     #[simd_test = "avx"]
     unsafe fn _mm256_extract_epi32() {
-        let a = i32x8::new(1, 2, 3, 4, 5, 6, 7, 8);
-        let r = avx::_mm256_extract_epi32(a, 0);
-        assert_eq!(r, 1);
+        let a = i32x8::new(-1, 1, 2, 3, 4, 5, 6, 7);
+        let r1 = avx::_mm256_extract_epi32(a, 0);
+        let r2 = avx::_mm256_extract_epi32(a, 11);
+        assert_eq!(r1, -1);
+        assert_eq!(r2, 3);
     }
 
     #[simd_test = "avx"]

coresimd/src/x86/i586/sse2.rs

@@ -971,7 +971,8 @@ pub unsafe fn _mm_packus_epi16(a: i16x8, b: i16x8) -> u8x16 {
 #[target_feature = "+sse2"]
 #[cfg_attr(test, assert_instr(pextrw, imm8 = 9))]
 pub unsafe fn _mm_extract_epi16(a: i16x8, imm8: i32) -> i32 {
-    a.extract(imm8 as u32 & 0b111) as i32
+    let imm8 = (imm8 & 7) as u32;
+    (a.extract_unchecked(imm8) as i32) & 0xFFFF
 }
 
 /// Return a new vector where the `imm8` element of `a` is replaced with `i`.
@@ -3122,9 +3123,11 @@ mod tests {
 
     #[simd_test = "sse2"]
     unsafe fn _mm_extract_epi16() {
-        let a = i16x8::new(0, 1, 2, 3, 4, 5, 6, 7);
-        let r = sse2::_mm_extract_epi16(a, 5);
-        assert_eq!(r, 5);
+        let a = i16x8::new(-1, 1, 2, 3, 4, 5, 6, 7);
+        let r1 = sse2::_mm_extract_epi16(a, 0);
+        let r2 = sse2::_mm_extract_epi16(a, 11);
+        assert_eq!(r1, 0xFFFF);
+        assert_eq!(r2, 3);
     }
 
     #[simd_test = "sse2"]

coresimd/src/x86/i586/sse41.rs

@@ -120,21 +120,25 @@ pub unsafe fn _mm_extract_ps(a: f32x4, imm8: u8) -> i32 {
     mem::transmute(a.extract(imm8 as u32 & 0b11))
 }
 
-/// Extract an 8-bit integer from `a` selected with `imm8`
+/// Extract an 8-bit integer from `a`, selected with `imm8`. Returns a 32-bit
+/// integer containing the zero-extended integer data.
+/// See: https://reviews.llvm.org/D20468
 #[inline(always)]
 #[target_feature = "+sse4.1"]
 #[cfg_attr(test, assert_instr(pextrb, imm8 = 0))]
-pub unsafe fn _mm_extract_epi8(a: i8x16, imm8: u8) -> i8 {
-    a.extract((imm8 & 0b1111) as u32)
+pub unsafe fn _mm_extract_epi8(a: i8x16, imm8: i32) -> i32 {
+    let imm8 = (imm8 & 15) as u32;
+    (a.extract_unchecked(imm8) as i32) & 0xFF
 }
 
 /// Extract an 32-bit integer from `a` selected with `imm8`
 #[inline(always)]
 #[target_feature = "+sse4.1"]
 // TODO: Add test for Windows
 #[cfg_attr(all(test, not(windows)), assert_instr(pextrd, imm8 = 1))]
-pub unsafe fn _mm_extract_epi32(a: i32x4, imm8: u8) -> i32 {
-    a.extract((imm8 & 0b11) as u32)
+pub unsafe fn _mm_extract_epi32(a: i32x4, imm8: i32) -> i32 {
+    let imm8 = (imm8 & 3) as u32;
+    a.extract_unchecked(imm8) as i32
 }
 
 /// Select a single value in `a` to store at some position in `b`,
@@ -844,11 +848,11 @@ mod tests {
     #[simd_test = "sse4.1"]
     unsafe fn _mm_extract_epi8() {
         let a =
-            i8x16::new(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
-        let r = sse41::_mm_extract_epi8(a, 1);
-        assert_eq!(r, 1);
-        let r = sse41::_mm_extract_epi8(a, 17);
-        assert_eq!(r, 1);
+            i8x16::new(-1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+        let r1 = sse41::_mm_extract_epi8(a, 0);
+        let r2 = sse41::_mm_extract_epi8(a, 19);
+        assert_eq!(r1, 0xFF);
+        assert_eq!(r2, 3);
     }
 
     #[simd_test = "sse4.1"]

coresimd/src/x86/i686/sse41.rs

@@ -24,8 +24,9 @@ extern "C" {
 // On x86 this emits 2 pextrd instructions
 #[cfg_attr(all(test, not(windows), target_arch = "x86"),
            assert_instr(pextrd, imm8 = 1))]
-pub unsafe fn _mm_extract_epi64(a: i64x2, imm8: u8) -> i64 {
-    a.extract((imm8 & 0b1) as u32)
+pub unsafe fn _mm_extract_epi64(a: i64x2, imm8: i32) -> i64 {
+    let imm8 = (imm8 & 1) as u32;
+    a.extract_unchecked(imm8)
 }
 
 /// Return a copy of `a` with the 64-bit integer from `i` inserted at a