Use target_feature annotation to mark functions requiring AVX2

zakcutner · marmeladema · commit 0583614e723c · 2020-07-21T11:24:35.000+01:00
diff --git a/benches/i386.rs b/benches/i386.rs
@@ -143,7 +143,7 @@ fn search_short_haystack(c: &mut Criterion) {
     group.bench_function("DynamicAvx2Searcher::search_in", |b| {
         let searchers = needles
             .iter()
-            .map(|&needle| {
+            .map(|&needle| unsafe {
                 DynamicAvx2Searcher::new(needle.as_bytes().to_owned().into_boxed_slice())
             })
             .collect::<Vec<_>>();
@@ -274,7 +274,9 @@ fn search_long_haystack(c: &mut Criterion) {
     group.bench_function("DynamicAvx2Searcher::search_in", |b| {
         let searchers = needles
             .iter()
-            .map(|needle| DynamicAvx2Searcher::new(needle.as_bytes().to_owned().into_boxed_slice()))
+            .map(|needle| unsafe {
+                DynamicAvx2Searcher::new(needle.as_bytes().to_owned().into_boxed_slice())
+            })
             .collect::<Vec<_>>();
 
         b.iter(|| {
diff --git a/benches/random.rs b/benches/random.rs
@@ -105,7 +105,8 @@ fn search(c: &mut Criterion) {
                 BenchmarkId::new("DynamicAvx2Searcher::search_in", parameter),
                 &size,
                 |b, _| {
-                    let searcher = DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice());
+                    let searcher =
+                        unsafe { DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice()) };
                     b.iter(|| black_box(searcher.search_in(haystack)));
                 },
             );
diff --git a/src/avx2/mod.rs b/src/avx2/mod.rs
@@ -1,3 +1,5 @@
+#![allow(clippy::missing_safety_doc)]
+
 mod original;
 mod rust;
 
@@ -45,53 +47,63 @@ trait Vector: Copy {
 
 impl Vector for __m128i {
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn set1_epi8(a: i8) -> Self {
         _mm_set1_epi8(a)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn loadu_si(a: *const Self) -> Self {
         _mm_loadu_si128(a)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn cmpeq_epi8(a: Self, b: Self) -> Self {
         _mm_cmpeq_epi8(a, b)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn and_si(a: Self, b: Self) -> Self {
         _mm_and_si128(a, b)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn movemask_epi8(a: Self) -> i32 {
         _mm_movemask_epi8(a)
     }
 }
 
 impl Vector for __m256i {
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn set1_epi8(a: i8) -> Self {
         _mm256_set1_epi8(a)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn loadu_si(a: *const Self) -> Self {
         _mm256_loadu_si256(a)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn cmpeq_epi8(a: Self, b: Self) -> Self {
         _mm256_cmpeq_epi8(a, b)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn and_si(a: Self, b: Self) -> Self {
         _mm256_and_si256(a, b)
     }
 
     #[inline]
+    #[target_feature(enable = "avx2")]
     unsafe fn movemask_epi8(a: Self) -> i32 {
         _mm256_movemask_epi8(a)
     }
@@ -103,10 +115,11 @@ struct VectorHash<V: Vector> {
 }
 
 impl<V: Vector> VectorHash<V> {
-    fn new(first: u8, last: u8) -> Self {
+    #[target_feature(enable = "avx2")]
+    unsafe fn new(first: u8, last: u8) -> Self {
         Self {
-            first: unsafe { Vector::set1_epi8(first as i8) },
-            last: unsafe { Vector::set1_epi8(last as i8) },
+            first: Vector::set1_epi8(first as i8),
+            last: Vector::set1_epi8(last as i8),
         }
     }
 }
@@ -122,12 +135,14 @@ macro_rules! avx2_searcher {
         }
 
         impl $name {
-            pub fn new(needle: Box<[u8]>) -> Self {
+            #[target_feature(enable = "avx2")]
+            pub unsafe fn new(needle: Box<[u8]>) -> Self {
                 let position = needle.len() - 1;
                 Self::with_position(needle, position)
             }
 
-            pub fn with_position(needle: Box<[u8]>, position: usize) -> Self {
+            #[target_feature(enable = "avx2")]
+            pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {
                 assert!(!needle.is_empty());
                 assert!(position < needle.len());
 
@@ -176,6 +191,7 @@ macro_rules! avx2_searcher {
             }
 
             #[inline]
+            #[target_feature(enable = "avx2")]
             unsafe fn vector_search_in_chunk<V: Vector>(
                 &self,
                 haystack: &[u8],
@@ -206,11 +222,12 @@ macro_rules! avx2_searcher {
             }
 
             #[inline]
-            fn vector_search_in<V: Vector>(
+            #[target_feature(enable = "avx2")]
+            unsafe fn vector_search_in<V: Vector>(
                 &self,
                 haystack: &[u8],
                 hash: &VectorHash<V>,
-                next: fn(&Self, &[u8]) -> bool,
+                next: unsafe fn(&Self, &[u8]) -> bool,
             ) -> bool {
                 debug_assert!(haystack.len() >= self.size());
 
@@ -223,17 +240,17 @@ macro_rules! avx2_searcher {
 
                 let mut chunks = haystack[..end].chunks_exact(lanes);
                 while let Some(chunk) = chunks.next() {
-                    if unsafe { self.vector_search_in_chunk(haystack, hash, chunk.as_ptr(), -1) } {
+                    if self.vector_search_in_chunk(haystack, hash, chunk.as_ptr(), -1) {
                         return true;
                     }
                 }
 
                 let remainder = chunks.remainder().len();
                 if remainder > 0 {
-                    let start = unsafe { haystack.as_ptr().add(end - lanes) };
+                    let start = haystack.as_ptr().add(end - lanes);
                     let mask = -1 << (lanes - remainder);
 
-                    if unsafe { self.vector_search_in_chunk(haystack, hash, start, mask) } {
+                    if self.vector_search_in_chunk(haystack, hash, start, mask) {
                         return true;
                     }
                 }
@@ -242,26 +259,30 @@ macro_rules! avx2_searcher {
             }
 
             #[inline]
-            fn sse2_search_in(&self, haystack: &[u8]) -> bool {
+            #[target_feature(enable = "avx2")]
+            unsafe fn sse2_search_in(&self, haystack: &[u8]) -> bool {
                 self.vector_search_in(haystack, &self.sse2_hash, Self::scalar_search_in)
             }
 
             #[inline]
-            fn avx2_search_in(&self, haystack: &[u8]) -> bool {
+            #[target_feature(enable = "avx2")]
+            unsafe fn avx2_search_in(&self, haystack: &[u8]) -> bool {
                 self.vector_search_in(haystack, &self.avx2_hash, Self::sse2_search_in)
             }
 
             #[inline]
-            pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {
+            #[target_feature(enable = "avx2")]
+            pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {
                 if haystack.len() < self.size() {
                     return false;
                 }
 
                 self.avx2_search_in(haystack)
             }
 
+            #[inline]
             pub fn search_in(&self, haystack: &[u8]) -> bool {
-                self.inlined_search_in(haystack)
+                unsafe { self.inlined_search_in(haystack) }
             }
         }
     };
@@ -300,12 +321,14 @@ pub enum DynamicAvx2Searcher {
 }
 
 impl DynamicAvx2Searcher {
-    pub fn new(needle: Box<[u8]>) -> Self {
+    #[target_feature(enable = "avx2")]
+    pub unsafe fn new(needle: Box<[u8]>) -> Self {
         let position = needle.len() - 1;
         Self::with_position(needle, position)
     }
 
-    pub fn with_position(needle: Box<[u8]>, position: usize) -> Self {
+    #[target_feature(enable = "avx2")]
+    pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {
         assert!(!needle.is_empty());
         assert!(position < needle.len());
 
@@ -329,7 +352,8 @@ impl DynamicAvx2Searcher {
     }
 
     #[inline]
-    pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {
+    #[target_feature(enable = "avx2")]
+    pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {
         match self {
             Self::N0 => true,
             Self::N1(searcher) => searcher.inlined_search_in(haystack),
@@ -349,8 +373,9 @@ impl DynamicAvx2Searcher {
         }
     }
 
+    #[inline]
     pub fn search_in(&self, haystack: &[u8]) -> bool {
-        self.inlined_search_in(haystack)
+        unsafe { self.inlined_search_in(haystack) }
     }
 }
 
@@ -359,7 +384,7 @@ mod tests {
     use super::Avx2Searcher;
 
     fn search(haystack: &[u8], needle: &[u8]) -> bool {
-        let search = |position| {
+        let search = |position| unsafe {
             Avx2Searcher::with_position(needle.to_owned().into_boxed_slice(), position)
                 .search_in(haystack)
         };
diff --git a/tests/i386.rs b/tests/i386.rs
@@ -37,7 +37,7 @@ fn search(haystack: &str, needle: &str) {
     let searcher = StrStrAVX2Searcher::new(needle);
     assert_eq!(searcher.search_in(haystack), result);
 
-    let searcher = DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice());
+    let searcher = unsafe { DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice()) };
     assert_eq!(searcher.search_in(haystack), result);
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#![allow(clippy::missing_safety_doc)]`
	`2`	`+`
`1`	`3`	`mod original;`
`2`	`4`	`mod rust;`
`3`	`5`
`@@ -45,53 +47,63 @@ trait Vector: Copy {`
`45`	`47`
`46`	`48`	`impl Vector for __m128i {`
`47`	`49`	`#[inline]`
	`50`	`+ #[target_feature(enable = "avx2")]`
`48`	`51`	`unsafe fn set1_epi8(a: i8) -> Self {`
`49`	`52`	`_mm_set1_epi8(a)`
`50`	`53`	`}`
`51`	`54`
`52`	`55`	`#[inline]`
	`56`	`+ #[target_feature(enable = "avx2")]`
`53`	`57`	`unsafe fn loadu_si(a: *const Self) -> Self {`
`54`	`58`	`_mm_loadu_si128(a)`
`55`	`59`	`}`
`56`	`60`
`57`	`61`	`#[inline]`
	`62`	`+ #[target_feature(enable = "avx2")]`
`58`	`63`	`unsafe fn cmpeq_epi8(a: Self, b: Self) -> Self {`
`59`	`64`	`_mm_cmpeq_epi8(a, b)`
`60`	`65`	`}`
`61`	`66`
`62`	`67`	`#[inline]`
	`68`	`+ #[target_feature(enable = "avx2")]`
`63`	`69`	`unsafe fn and_si(a: Self, b: Self) -> Self {`
`64`	`70`	`_mm_and_si128(a, b)`
`65`	`71`	`}`
`66`	`72`
`67`	`73`	`#[inline]`
	`74`	`+ #[target_feature(enable = "avx2")]`
`68`	`75`	`unsafe fn movemask_epi8(a: Self) -> i32 {`
`69`	`76`	`_mm_movemask_epi8(a)`
`70`	`77`	`}`
`71`	`78`	`}`
`72`	`79`
`73`	`80`	`impl Vector for __m256i {`
`74`	`81`	`#[inline]`
	`82`	`+ #[target_feature(enable = "avx2")]`
`75`	`83`	`unsafe fn set1_epi8(a: i8) -> Self {`
`76`	`84`	`_mm256_set1_epi8(a)`
`77`	`85`	`}`
`78`	`86`
`79`	`87`	`#[inline]`
	`88`	`+ #[target_feature(enable = "avx2")]`
`80`	`89`	`unsafe fn loadu_si(a: *const Self) -> Self {`
`81`	`90`	`_mm256_loadu_si256(a)`
`82`	`91`	`}`
`83`	`92`
`84`	`93`	`#[inline]`
	`94`	`+ #[target_feature(enable = "avx2")]`
`85`	`95`	`unsafe fn cmpeq_epi8(a: Self, b: Self) -> Self {`
`86`	`96`	`_mm256_cmpeq_epi8(a, b)`
`87`	`97`	`}`
`88`	`98`
`89`	`99`	`#[inline]`
	`100`	`+ #[target_feature(enable = "avx2")]`
`90`	`101`	`unsafe fn and_si(a: Self, b: Self) -> Self {`
`91`	`102`	`_mm256_and_si256(a, b)`
`92`	`103`	`}`
`93`	`104`
`94`	`105`	`#[inline]`
	`106`	`+ #[target_feature(enable = "avx2")]`
`95`	`107`	`unsafe fn movemask_epi8(a: Self) -> i32 {`
`96`	`108`	`_mm256_movemask_epi8(a)`
`97`	`109`	`}`
`@@ -103,10 +115,11 @@ struct VectorHash<V: Vector> {`
`103`	`115`	`}`
`104`	`116`
`105`	`117`	`impl<V: Vector> VectorHash<V> {`
`106`		`- fn new(first: u8, last: u8) -> Self {`
	`118`	`+ #[target_feature(enable = "avx2")]`
	`119`	`+ unsafe fn new(first: u8, last: u8) -> Self {`
`107`	`120`	`Self {`
`108`		`- first: unsafe { Vector::set1_epi8(first as i8) },`
`109`		`- last: unsafe { Vector::set1_epi8(last as i8) },`
	`121`	`+ first: Vector::set1_epi8(first as i8),`
	`122`	`+ last: Vector::set1_epi8(last as i8),`
`110`	`123`	`}`
`111`	`124`	`}`
`112`	`125`	`}`
`@@ -122,12 +135,14 @@ macro_rules! avx2_searcher {`
`122`	`135`	`}`
`123`	`136`
`124`	`137`	`impl $name {`
`125`		`- pub fn new(needle: Box<[u8]>) -> Self {`
	`138`	`+ #[target_feature(enable = "avx2")]`
	`139`	`+ pub unsafe fn new(needle: Box<[u8]>) -> Self {`
`126`	`140`	`let position = needle.len() - 1;`
`127`	`141`	`Self::with_position(needle, position)`
`128`	`142`	`}`
`129`	`143`
`130`		`- pub fn with_position(needle: Box<[u8]>, position: usize) -> Self {`
	`144`	`+ #[target_feature(enable = "avx2")]`
	`145`	`+ pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {`
`131`	`146`	`assert!(!needle.is_empty());`
`132`	`147`	`assert!(position < needle.len());`
`133`	`148`
`@@ -176,6 +191,7 @@ macro_rules! avx2_searcher {`
`176`	`191`	`}`
`177`	`192`
`178`	`193`	`#[inline]`
	`194`	`+ #[target_feature(enable = "avx2")]`
`179`	`195`	`unsafe fn vector_search_in_chunk<V: Vector>(`
`180`	`196`	`&self,`
`181`	`197`	`haystack: &[u8],`
`@@ -206,11 +222,12 @@ macro_rules! avx2_searcher {`
`206`	`222`	`}`
`207`	`223`
`208`	`224`	`#[inline]`
`209`		`- fn vector_search_in<V: Vector>(`
	`225`	`+ #[target_feature(enable = "avx2")]`
	`226`	`+ unsafe fn vector_search_in<V: Vector>(`
`210`	`227`	`&self,`
`211`	`228`	`haystack: &[u8],`
`212`	`229`	`hash: &VectorHash<V>,`
`213`		`- next: fn(&Self, &[u8]) -> bool,`
	`230`	`+ next: unsafe fn(&Self, &[u8]) -> bool,`
`214`	`231`	`) -> bool {`
`215`	`232`	`debug_assert!(haystack.len() >= self.size());`
`216`	`233`
`@@ -223,17 +240,17 @@ macro_rules! avx2_searcher {`
`223`	`240`
`224`	`241`	`let mut chunks = haystack[..end].chunks_exact(lanes);`
`225`	`242`	`while let Some(chunk) = chunks.next() {`
`226`		`- if unsafe { self.vector_search_in_chunk(haystack, hash, chunk.as_ptr(), -1) } {`
	`243`	`+ if self.vector_search_in_chunk(haystack, hash, chunk.as_ptr(), -1) {`
`227`	`244`	`return true;`
`228`	`245`	`}`
`229`	`246`	`}`
`230`	`247`
`231`	`248`	`let remainder = chunks.remainder().len();`
`232`	`249`	`if remainder > 0 {`
`233`		`- let start = unsafe { haystack.as_ptr().add(end - lanes) };`
	`250`	`+ let start = haystack.as_ptr().add(end - lanes);`
`234`	`251`	`let mask = -1 << (lanes - remainder);`
`235`	`252`
`236`		`- if unsafe { self.vector_search_in_chunk(haystack, hash, start, mask) } {`
	`253`	`+ if self.vector_search_in_chunk(haystack, hash, start, mask) {`
`237`	`254`	`return true;`
`238`	`255`	`}`
`239`	`256`	`}`
`@@ -242,26 +259,30 @@ macro_rules! avx2_searcher {`
`242`	`259`	`}`
`243`	`260`
`244`	`261`	`#[inline]`
`245`		`- fn sse2_search_in(&self, haystack: &[u8]) -> bool {`
	`262`	`+ #[target_feature(enable = "avx2")]`
	`263`	`+ unsafe fn sse2_search_in(&self, haystack: &[u8]) -> bool {`
`246`	`264`	`self.vector_search_in(haystack, &self.sse2_hash, Self::scalar_search_in)`
`247`	`265`	`}`
`248`	`266`
`249`	`267`	`#[inline]`
`250`		`- fn avx2_search_in(&self, haystack: &[u8]) -> bool {`
	`268`	`+ #[target_feature(enable = "avx2")]`
	`269`	`+ unsafe fn avx2_search_in(&self, haystack: &[u8]) -> bool {`
`251`	`270`	`self.vector_search_in(haystack, &self.avx2_hash, Self::sse2_search_in)`
`252`	`271`	`}`
`253`	`272`
`254`	`273`	`#[inline]`
`255`		`- pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {`
	`274`	`+ #[target_feature(enable = "avx2")]`
	`275`	`+ pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {`
`256`	`276`	`if haystack.len() < self.size() {`
`257`	`277`	`return false;`
`258`	`278`	`}`
`259`	`279`
`260`	`280`	`self.avx2_search_in(haystack)`
`261`	`281`	`}`
`262`	`282`
	`283`	`+ #[inline]`
`263`	`284`	`pub fn search_in(&self, haystack: &[u8]) -> bool {`
`264`		`- self.inlined_search_in(haystack)`
	`285`	`+ unsafe { self.inlined_search_in(haystack) }`
`265`	`286`	`}`
`266`	`287`	`}`
`267`	`288`	`};`
`@@ -300,12 +321,14 @@ pub enum DynamicAvx2Searcher {`
`300`	`321`	`}`
`301`	`322`
`302`	`323`	`impl DynamicAvx2Searcher {`
`303`		`- pub fn new(needle: Box<[u8]>) -> Self {`
	`324`	`+ #[target_feature(enable = "avx2")]`
	`325`	`+ pub unsafe fn new(needle: Box<[u8]>) -> Self {`
`304`	`326`	`let position = needle.len() - 1;`
`305`	`327`	`Self::with_position(needle, position)`
`306`	`328`	`}`
`307`	`329`
`308`		`- pub fn with_position(needle: Box<[u8]>, position: usize) -> Self {`
	`330`	`+ #[target_feature(enable = "avx2")]`
	`331`	`+ pub unsafe fn with_position(needle: Box<[u8]>, position: usize) -> Self {`
`309`	`332`	`assert!(!needle.is_empty());`
`310`	`333`	`assert!(position < needle.len());`
`311`	`334`
`@@ -329,7 +352,8 @@ impl DynamicAvx2Searcher {`
`329`	`352`	`}`
`330`	`353`
`331`	`354`	`#[inline]`
`332`		`- pub fn inlined_search_in(&self, haystack: &[u8]) -> bool {`
	`355`	`+ #[target_feature(enable = "avx2")]`
	`356`	`+ pub unsafe fn inlined_search_in(&self, haystack: &[u8]) -> bool {`
`333`	`357`	`match self {`
`334`	`358`	`Self::N0 => true,`
`335`	`359`	`Self::N1(searcher) => searcher.inlined_search_in(haystack),`
`@@ -349,8 +373,9 @@ impl DynamicAvx2Searcher {`
`349`	`373`	`}`
`350`	`374`	`}`
`351`	`375`
	`376`	`+ #[inline]`
`352`	`377`	`pub fn search_in(&self, haystack: &[u8]) -> bool {`
`353`		`- self.inlined_search_in(haystack)`
	`378`	`+ unsafe { self.inlined_search_in(haystack) }`
`354`	`379`	`}`
`355`	`380`	`}`
`356`	`381`
`@@ -359,7 +384,7 @@ mod tests {`
`359`	`384`	`use super::Avx2Searcher;`
`360`	`385`
`361`	`386`	`fn search(haystack: &[u8], needle: &[u8]) -> bool {`
`362`		`- let search = \|position\| {`
	`387`	`+ let search = \|position\| unsafe {`
`363`	`388`	`Avx2Searcher::with_position(needle.to_owned().into_boxed_slice(), position)`
`364`	`389`	`.search_in(haystack)`
`365`	`390`	`};`
Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ fn search(haystack: &str, needle: &str) {`
`37`	`37`	`let searcher = StrStrAVX2Searcher::new(needle);`
`38`	`38`	`assert_eq!(searcher.search_in(haystack), result);`
`39`	`39`
`40`		`- let searcher = DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice());`
	`40`	`+ let searcher = unsafe { DynamicAvx2Searcher::new(needle.to_owned().into_boxed_slice()) };`
`41`	`41`	`assert_eq!(searcher.search_in(haystack), result);`
`42`	`42`	`}`
`43`	`43`