Hot-path acceleration: one-pass scan and vectorization

Author: hghwng

Makefile

@@ -27,7 +27,7 @@ MISC_PATH   = $(PREFIX)/share/afl
 PROGS       = afl-gcc afl-fuzz afl-showmap afl-tmin afl-gotcpu afl-analyze
 SH_PROGS    = afl-plot afl-cmin afl-whatsup
 
-CFLAGS     ?= -O3 -funroll-loops
+CFLAGS     ?= -O3 -march=native
 CFLAGS     += -Wall -D_FORTIFY_SOURCE=2 -g -Wno-pointer-sign \
 	      -DAFL_PATH=\"$(HELPER_PATH)\" -DDOC_PATH=\"$(DOC_PATH)\" \
 	      -DBIN_PATH=\"$(BIN_PATH)\"

afl-fuzz.c

@@ -1096,7 +1096,21 @@ static inline u8 has_new_bits(u8* virgin_map) {
  * return has_new_bits(). */
 
 static inline u8 has_new_bits_unclassified(u8* virgin_map) {
-  classify_counts(trace_bits); // TODO
+
+  /* Handle the hot path first: no new coverage */
+  u8* end = trace_bits + MAP_SIZE;
+
+#ifdef WORD_SIZE_64
+
+  if (!skim((u64*)virgin_map, (u64*)trace_bits, (u64*)end)) return 0;
+
+#else
+
+  if (!skim((u32*)virgin_map, (u32*)trace_bits, (u32*)end)) return 0;
+
+#endif /* ^WORD_SIZE_64 */
+
+  classify_counts(trace_bits);
   return has_new_bits(virgin_map);
 }
 
@@ -3044,7 +3058,19 @@ static u8 save_if_interesting(char** argv, void* mem, u32 len, u8 fault) {
     /* Keep only if there are new bits in the map, add to queue for
        future fuzzing, etc. */
 
-    if (!(hnb = has_new_bits_unclassified(virgin_bits))) {
+
+    /* A combination of classify_counts and has_new_bits. If 0 is returned, then
+     * the trace bits are kept as-is. Otherwise, the trace bits are overwritten
+     * with classified values.
+     *
+     * This accelerates the processing: in most cases, no interesting behavior
+     * happen, and the trace bits will be discarded soon. This function
+     * optimizes for such cases: one-pass scan on trace bits without modifying
+     * anything. Only on rare cases it fall backs to the slow path:
+     * classify_counts() first, then return has_new_bits(). */
+    hnb = has_new_bits_unclassified(virgin_bits);
+
+    if (!hnb) {
       if (crash_mode) total_crashes++;
       return 0;
     }    

coverage-32.h

@@ -84,3 +84,19 @@ static inline void discover_word(u8* ret, u32* current, u32* virgin) {
     *virgin &= ~*current;
   }
 }
+
+
+#define PACK_SIZE 16
+static inline u32 skim(const u32* virgin, const u32* current, const u32* current_end) {
+
+  for (; current != current_end; virgin += 4, current += 4) {
+
+    if (current[0] && classify_word(current[0]) & virgin[0]) return 1;
+    if (current[1] && classify_word(current[1]) & virgin[1]) return 1;
+    if (current[2] && classify_word(current[2]) & virgin[2]) return 1;
+    if (current[3] && classify_word(current[3]) & virgin[3]) return 1;
+
+  }
+
+  return 0;
+}

coverage-64.h

@@ -1,6 +1,9 @@
 #include "config.h"
 #include "types.h"
 
+#if (defined(__AVX512F__) && defined(__AVX512DQ__)) || defined(__AVX2__)
+#  include <immintrin.h>
+#endif
 
 static inline u64 classify_word(u64 word) {
 
@@ -94,3 +97,78 @@ static inline void discover_word(u8* ret, u64* current, u64* virgin) {
   }
 
 }
+
+
+#if defined(__AVX512F__) && defined(__AVX512DQ__)
+#define PACK_SIZE 64
+static inline u32 skim(const u64* virgin, const u64* current, const u64* current_end) {
+
+  for (; current != current_end; virgin += 8, current += 8) {
+
+    __m512i value = *(__m512i*)current;
+    __mmask8 mask = _mm512_testn_epi64_mask(value, value);
+
+    /* All bytes are zero. */
+    if (mask == 0xff) continue;
+
+    /* Look for nonzero bytes and check for new bits. */
+#define UNROLL(x) \
+  if (!(mask & (1 << x)) && classify_word(current[x]) & virgin[x]) return 1
+    UNROLL(0); UNROLL(1); UNROLL(2); UNROLL(3);
+    UNROLL(4); UNROLL(5); UNROLL(6); UNROLL(7);
+#undef UNROLL
+
+  }
+
+  return 0;
+
+}
+#endif
+
+
+#if !defined(PACK_SIZE) && defined(__AVX2__)
+#define PACK_SIZE 32
+static inline u32 skim(const u64* virgin, const u64* current, const u64* current_end) {
+
+  __m256i zeroes = _mm256_setzero_si256();
+
+  for (; current != current_end; virgin += 4, current += 4) {
+
+    __m256i value = *(__m256i*)current;
+    __m256i cmp = _mm256_cmpeq_epi64(value, zeroes);
+    u32 mask = _mm256_movemask_epi8(cmp);
+
+    /* All bytes are zero. */
+    if (mask == -1) continue;
+
+    /* Look for nonzero bytes and check for new bits. */
+    if (!(mask & 0xff) && classify_word(current[0]) & virgin[0]) return 1;
+    if (!(mask & 0xff00) && classify_word(current[1]) & virgin[1]) return 1;
+    if (!(mask & 0xff0000) && classify_word(current[2]) & virgin[2]) return 1;
+    if (!(mask & 0xff000000) && classify_word(current[3]) & virgin[3]) return 1;
+
+  }
+
+  return 0;
+
+}
+#endif
+
+
+#if !defined(PACK_SIZE)
+#define PACK_SIZE 32
+static inline u32 skim(const u64* virgin, const u64* current, const u64* current_end) {
+
+  for (; current != current_end; virgin += 4, current += 4) {
+
+    if (current[0] && classify_word(current[0]) & virgin[0]) return 1;
+    if (current[1] && classify_word(current[1]) & virgin[1]) return 1;
+    if (current[2] && classify_word(current[2]) & virgin[2]) return 1;
+    if (current[3] && classify_word(current[3]) & virgin[3]) return 1;
+
+  }
+
+  return 0;
+
+}
+#endif