Skip to content

Commit bfc55c9

Browse files
wtdcodetokatoka
wtdcode
and
tokatoka
authored
Enable SIMD acceleration for stable rust toolchain (#3140)
* initial support * migrate SAND * Update comments * Fmt * Clippy * Fix missing docs * fmt fix * clippy again * weird clippy * clippy * Fix * Allow new SIMDMapFeedback * Fix features * Fix features again * Allow custom names * Fix imports * Fix imports * Fmt * Fix missing implementations * Requires std to simd * DO NOT Overwrite names * Format toml * no_std fix * fmt * Use SIMDMapFeedback for libfuzzer_libpng * no_std (?) * clippy * fix no_alloc * allow cargo docs to enable all features * clippy again * Fix missing import * Fix cargo docs * Naive simplify_map doesn't require wide * Accidentally commit the file * more fine grined features * Fix clippy.ps1 * Fix wide256 for simplify_map * Renaming to SimdMapFeedback * Dynamic dispatch * Fix naming * Move to simd.rs * clippy * clippy * dispatch earlier * Fix clippy * clippy * clippy * Revert previous change * Fix comments * Update comments for std_covmap_is_interesting * remove SIMD and choose fastest implementation based on target_arch * no longer nightly imports * Fix * upstream benchmark code * Fix docs * Fix libfuzzer_libpng * Disable clippy for benchmark * clippy * clippy again --------- Co-authored-by: Dongjia "toka" Zhang <[email protected]>
1 parent 373fe03 commit bfc55c9

File tree

10 files changed

+791
-123
lines changed

10 files changed

+791
-123
lines changed

‎fuzzers/inprocess/libfuzzer_libpng/src/lib.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ pub extern "C" fn libafl_main() {
6262
#[cfg(not(test))]
6363
fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Result<(), Error> {
6464
// 'While the stats are state, they are usually used in the broker - which is likely never restarted
65+
66+
use libafl::feedbacks::simd::{SimdImplmentation, SimdMapFeedback};
6567
let monitor = MultiMonitor::new(|s| println!("{s}"));
6668

6769
// The restarting state will spawn the same process again as child, then restarted it each time it crashes.
@@ -93,8 +95,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
9395
// Create an observation channel to keep track of the execution time
9496
let time_observer = TimeObserver::new("time");
9597

96-
let map_feedback = MaxMapFeedback::new(&edges_observer);
97-
98+
let map_feedback = SimdMapFeedback::new(MaxMapFeedback::new(&edges_observer));
9899
let calibration = CalibrationStage::new(&map_feedback);
99100

100101
// Feedback to rate the interestingness of an input

‎libafl/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ default = [
3939
"regex",
4040
"serdeany_autoreg",
4141
"libafl_bolts/xxh3",
42+
"stable_simd",
4243
]
4344
document-features = ["dep:document-features"]
4445

@@ -195,6 +196,9 @@ nautilus = [
195196
"regex",
196197
]
197198

199+
## Use the best SIMD implementation by our [benchmark](https://github.com/wtdcode/libafl_simd_bench)
200+
stable_simd = ["libafl_bolts/stable_simd"]
201+
198202
[[example]]
199203
name = "tui_mock"
200204
path = "./examples/tui_mock/main.rs"

‎libafl/src/executors/sand.rs

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use core::marker::PhantomData;
99

1010
use libafl_bolts::{
1111
AsIter, Error, Named, hash_std,
12+
simd::std_simplify_map,
1213
tuples::{Handle, MatchName, MatchNameRef},
1314
};
1415

@@ -148,14 +149,10 @@ where
148149
let kind = self.executor.run_target(fuzzer, state, mgr, input)?;
149150
let ot = self.executor.observers();
150151
let ob = ot.get(&self.ob_ref).unwrap().as_ref();
151-
let initial = ob.initial();
152152
let mut covs = ob.to_vec();
153153
match self.pattern {
154154
SANDExecutionPattern::SimplifiedTrace => {
155-
// TODO: SIMD Optimizations
156-
for it in &mut covs {
157-
*it = if *it == initial { 0x1 } else { 0x80 };
158-
}
155+
std_simplify_map(&mut covs);
159156
}
160157
SANDExecutionPattern::UniqueTrace => {
161158
classify_counts(covs.as_mut_slice());

‎libafl/src/feedbacks/map.rs

Lines changed: 64 additions & 116 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,16 @@
11
//! Map feedback, maximizing or minimizing maps, for example the afl-style map observer.
22
33
use alloc::{borrow::Cow, vec::Vec};
4-
#[rustversion::nightly]
5-
use core::simd::prelude::SimdOrd;
64
use core::{
75
fmt::Debug,
86
marker::PhantomData,
97
ops::{BitAnd, BitOr, Deref, DerefMut},
108
};
119

1210
#[rustversion::nightly]
13-
use libafl_bolts::AsSlice;
11+
use libafl_bolts::simd::std_covmap_is_interesting;
1412
use libafl_bolts::{
15-
AsIter, HasRefCnt, Named,
13+
AsIter, AsSlice, HasRefCnt, Named,
1614
tuples::{Handle, Handled, MatchName, MatchNameRef},
1715
};
1816
use num_traits::PrimInt;
@@ -548,7 +546,7 @@ where
548546
observers: &OT,
549547
_exit_kind: &ExitKind,
550548
) -> Result<bool, Error> {
551-
Ok(self.is_interesting_u8_simd_optimized(state, observers))
549+
Ok(self.is_interesting_u8_simd_optimized(state, observers, std_covmap_is_interesting))
552550
}
553551
}
554552

@@ -604,117 +602,6 @@ where
604602
}
605603
}
606604

607-
/// Specialize for the common coverage map size, maximization of u8s
608-
#[rustversion::nightly]
609-
impl<C, O> MapFeedback<C, DifferentIsNovel, O, MaxReducer>
610-
where
611-
O: MapObserver<Entry = u8> + for<'a> AsSlice<'a, Entry = u8> + for<'a> AsIter<'a, Item = u8>,
612-
C: CanTrack + AsRef<O>,
613-
{
614-
fn is_interesting_u8_simd_optimized<S, OT>(&mut self, state: &mut S, observers: &OT) -> bool
615-
where
616-
S: HasNamedMetadata,
617-
OT: MatchName,
618-
{
619-
// 128 bits vectors
620-
type VectorType = core::simd::u8x16;
621-
622-
let mut interesting = false;
623-
// TODO Replace with match_name_type when stable
624-
let observer = observers.get(&self.map_ref).expect("MapObserver not found. This is likely because you entered the crash handler with the wrong executor/observer").as_ref();
625-
626-
let map_state = state
627-
.named_metadata_map_mut()
628-
.get_mut::<MapFeedbackMetadata<u8>>(&self.name)
629-
.unwrap();
630-
let size = observer.usable_count();
631-
let len = observer.len();
632-
if map_state.history_map.len() < len {
633-
map_state.history_map.resize(len, u8::default());
634-
}
635-
636-
let map = observer.as_slice();
637-
debug_assert!(map.len() >= size);
638-
639-
let history_map = map_state.history_map.as_slice();
640-
641-
// Non vector implementation for reference
642-
/*for (i, history) in history_map.iter_mut().enumerate() {
643-
let item = map[i];
644-
let reduced = MaxReducer::reduce(*history, item);
645-
if DifferentIsNovel::is_novel(*history, reduced) {
646-
*history = reduced;
647-
interesting = true;
648-
if self.novelties.is_some() {
649-
self.novelties.as_mut().unwrap().push(i);
650-
}
651-
}
652-
}*/
653-
654-
let steps = size / VectorType::LEN;
655-
let left = size % VectorType::LEN;
656-
657-
if let Some(novelties) = self.novelties.as_mut() {
658-
novelties.clear();
659-
for step in 0..steps {
660-
let i = step * VectorType::LEN;
661-
let history = VectorType::from_slice(&history_map[i..]);
662-
let items = VectorType::from_slice(&map[i..]);
663-
664-
if items.simd_max(history) != history {
665-
interesting = true;
666-
unsafe {
667-
for j in i..(i + VectorType::LEN) {
668-
let item = *map.get_unchecked(j);
669-
if item > *history_map.get_unchecked(j) {
670-
novelties.push(j);
671-
}
672-
}
673-
}
674-
}
675-
}
676-
677-
for j in (size - left)..size {
678-
unsafe {
679-
let item = *map.get_unchecked(j);
680-
if item > *history_map.get_unchecked(j) {
681-
interesting = true;
682-
novelties.push(j);
683-
}
684-
}
685-
}
686-
} else {
687-
for step in 0..steps {
688-
let i = step * VectorType::LEN;
689-
let history = VectorType::from_slice(&history_map[i..]);
690-
let items = VectorType::from_slice(&map[i..]);
691-
692-
if items.simd_max(history) != history {
693-
interesting = true;
694-
break;
695-
}
696-
}
697-
698-
if !interesting {
699-
for j in (size - left)..size {
700-
unsafe {
701-
let item = *map.get_unchecked(j);
702-
if item > *history_map.get_unchecked(j) {
703-
interesting = true;
704-
break;
705-
}
706-
}
707-
}
708-
}
709-
}
710-
#[cfg(feature = "track_hit_feedbacks")]
711-
{
712-
self.last_result = Some(interesting);
713-
}
714-
interesting
715-
}
716-
}
717-
718605
impl<C, N, O, R> HasObserverHandle for MapFeedback<C, N, O, R> {
719606
type Observer = C;
720607

@@ -789,6 +676,67 @@ where
789676
}
790677
}
791678

679+
/// Specialize for the common coverage map size, maximization of u8s
680+
impl<C, O> MapFeedback<C, DifferentIsNovel, O, MaxReducer>
681+
where
682+
O: MapObserver<Entry = u8> + for<'a> AsSlice<'a, Entry = u8> + for<'a> AsIter<'a, Item = u8>,
683+
C: CanTrack + AsRef<O>,
684+
{
685+
#[allow(dead_code)] // this is true on stable wihout "stable_simd"
686+
pub(crate) fn is_interesting_u8_simd_optimized<S, OT, F>(
687+
&mut self,
688+
state: &mut S,
689+
observers: &OT,
690+
simd: F,
691+
) -> bool
692+
where
693+
S: HasNamedMetadata,
694+
OT: MatchName,
695+
F: FnOnce(&[u8], &[u8], bool) -> (bool, Vec<usize>),
696+
{
697+
// TODO Replace with match_name_type when stable
698+
let observer = observers.get(&self.map_ref).expect("MapObserver not found. This is likely because you entered the crash handler with the wrong executor/observer").as_ref();
699+
700+
let map_state = state
701+
.named_metadata_map_mut()
702+
.get_mut::<MapFeedbackMetadata<u8>>(&self.name)
703+
.unwrap();
704+
let size = observer.usable_count();
705+
let len = observer.len();
706+
if map_state.history_map.len() < len {
707+
map_state.history_map.resize(len, u8::default());
708+
}
709+
710+
let map = observer.as_slice();
711+
debug_assert!(map.len() >= size);
712+
713+
let history_map = map_state.history_map.as_slice();
714+
715+
// Non vector implementation for reference
716+
/*for (i, history) in history_map.iter_mut().enumerate() {
717+
let item = map[i];
718+
let reduced = MaxReducer::reduce(*history, item);
719+
if DifferentIsNovel::is_novel(*history, reduced) {
720+
*history = reduced;
721+
interesting = true;
722+
if self.novelties.is_some() {
723+
self.novelties.as_mut().unwrap().push(i);
724+
}
725+
}
726+
}*/
727+
728+
let (interesting, novelties) = simd(history_map, &map, self.novelties.is_some());
729+
if let Some(nov) = self.novelties.as_mut() {
730+
*nov = novelties;
731+
}
732+
#[cfg(feature = "track_hit_feedbacks")]
733+
{
734+
self.last_result = Some(interesting);
735+
}
736+
interesting
737+
}
738+
}
739+
792740
#[cfg(test)]
793741
mod tests {
794742
use crate::feedbacks::{AllIsNovel, IsNovel, NextPow2IsNovel};

‎libafl/src/feedbacks/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ pub mod map;
4646
pub mod nautilus;
4747
#[cfg(feature = "std")]
4848
pub mod new_hash_feedback;
49+
#[cfg(feature = "stable_simd")]
50+
pub mod simd;
4951
#[cfg(feature = "std")]
5052
pub mod stdio;
5153
pub mod transferred;

0 commit comments

Comments
 (0)