Implement detection and compilation of AVX512 code for the connection…

… scoring

CMakeLists.txt

@@ -9,6 +9,7 @@ set_property(GLOBAL PROPERTY PYTHON_EXTENSIONS_SOURCE_DIR ${CMAKE_CURRENT_SOURCE
 # --- Detect SIMD --------------------------------------------------------------
 
 include("src/scripts/cmake/FindAVX2.cmake")
+include("src/scripts/cmake/FindAVX512.cmake")
 include("src/scripts/cmake/FindNEON.cmake")
 include("src/scripts/cmake/FindSSE2.cmake")
 include("src/scripts/cmake/CythonExtension.cmake")

src/pyrodigal/impl/CMakeLists.txt

@@ -1,54 +1,33 @@
-# if(HAVE_SSE2)
-#   add_library(implsse sse.c sse.h template.h)
-#   target_include_directories(implsse PUBLIC $<TARGET_PROPERTY:prodigal,INCLUDE_DIRECTORIES>)
-#   if(SSE2_C_FLAGS AND NOT SSE2_C_FLAGS STREQUAL " ")
-#     target_link_options(implsse PUBLIC ${SSE2_C_FLAGS})
-#     target_compile_options(implsse PUBLIC ${SSE2_C_FLAGS})
-#   endif()
-# endif()
-
-# if(HAVE_AVX2)
-#   add_library(implavx avx.c avx.h template.h)
-#   target_include_directories(implavx PUBLIC  $<TARGET_PROPERTY:prodigal,INCLUDE_DIRECTORIES>)
-#   if(AVX2_C_FLAGS AND NOT AVX2_C_FLAGS STREQUAL " ")
-#     target_link_options(implavx PUBLIC ${AVX2_C_FLAGS})
-#     target_compile_options(implavx PUBLIC ${AVX2_C_FLAGS})
-#   endif()
-# endif()
-
-# if(HAVE_NEON)
-#   add_library(implneon neon.c neon.h template.h)
-#   target_include_directories(implneon PUBLIC  $<TARGET_PROPERTY:prodigal,INCLUDE_DIRECTORIES>)
-#   if(NEON_C_FLAGS AND NOT NEON_C_FLAGS STREQUAL " ")
-#     target_link_options(implneon PUBLIC ${NEON_C_FLAGS})
-#     target_compile_options(implneon PUBLIC ${NEON_C_FLAGS})
-#   endif()
-# endif()
-
-# add_library(impl generic.c generic.h template.h)
-# target_include_directories(impl PUBLIC $<TARGET_PROPERTY:prodigal,INCLUDE_DIRECTORIES> ${PROJECT_SOURCE_DIR}/src/pyrodigal/prodigal)
-
-
-
 cython_extension(generic LINKS prodigal EXTRA_SOURCES generic.c generic.h template.h)
+
 if(HAVE_SSE2)
   cython_extension(sse2 LINKS prodigal EXTRA_SOURCES sse2.c sse2.h template.h)
-  string(STRIP "${SSE2_C_FLAGS}" IMPL_FLAGS)
-  if(IMPL_FLAGS)
-    target_compile_options(pyrodigal.impl.sse2 PUBLIC ${IMPL_FLAGS})
-  endif()
+  string(REPLACE " " ";" IMPL_FLAGS ${SSE2_C_FLAGS})
+  foreach(_flag IN LISTS IMPL_FLAGS) 
+    target_compile_options(pyrodigal.impl.sse2 PUBLIC ${_flag})
+  endforeach()
 endif()
+
 if(HAVE_AVX2)
   cython_extension(avx2 LINKS prodigal EXTRA_SOURCES avx2.c avx2.h template.h)
-  string(STRIP "${AVX2_C_FLAGS}" IMPL_FLAGS)
-  if(IMPL_FLAGS)
-    target_compile_options(pyrodigal.impl.avx2 PUBLIC ${IMPL_FLAGS})
-  endif()
+  string(REPLACE " " ";" IMPL_FLAGS ${AVX2_C_FLAGS})
+  foreach(_flag IN LISTS IMPL_FLAGS) 
+    target_compile_options(pyrodigal.impl.avx2 PUBLIC ${_flag})
+  endforeach()
 endif()
+
+if(HAVE_AVX512)
+  cython_extension(avx512 LINKS prodigal EXTRA_SOURCES avx512.c avx512.h template.h)
+  string(REPLACE " " ";" IMPL_FLAGS ${AVX512_C_FLAGS})
+  foreach(_flag IN LISTS IMPL_FLAGS) 
+    target_compile_options(pyrodigal.impl.avx512 PUBLIC ${_flag})
+  endforeach()
+endif()
+
 if(HAVE_NEON)
   cython_extension(neon LINKS prodigal EXTRA_SOURCES neon.c neon.h template.h)
-  string(STRIP "${NEON_C_FLAGS}" IMPL_FLAGS)
-  if(IMPL_FLAGS)
-    target_compile_options(pyrodigal.impl.neon PUBLIC ${IMPL_FLAGS})
-  endif()
+  string(REPLACE " " ";" IMPL_FLAGS ${NEON_C_FLAGS})
+  foreach(_flag IN LISTS IMPL_FLAGS) 
+    target_compile_options(pyrodigal.impl.neon PUBLIC ${_flag})
+  endforeach()
 endif()

src/pyrodigal/impl/avx512.h

@@ -1,14 +1,39 @@
-#ifndef _PYRODIGAL_IMPL_AVX512_H
-#define _PYRODIGAL_IMPL_AVX512_H
+#include "sequence.h"
+#include "generic.h"
 
-#ifdef WIN32
-#define export __declspec( dllexport )
-#else
-#define export extern
-#endif
+#ifdef __AVX512F__
+#ifdef __AVX512BW__
+
+#include <immintrin.h>
+
+#include "template.h"
 
-#include <stdint.h>
+#define simd_t            __m512i
+#define simd_load(m)      _mm512_load_si512((__m512i*) (m))
+#define simd_store(x, m)  _mm512_store_si512((__m512i*) (m), x)
+#define simd_set1(x)      _mm512_set1_epi8(x)
+#define simd_eq(x, y)     _mm512_cmpeq_epi8_mask(x, y)
 
-export void skippable_avx512(const int8_t*, const uint8_t*, const uint8_t*, const int, const int, uint8_t*);
+#define mask_t            __mmask64
+#define mask_eq(x, y)     _kxnor_mask64(x, y)
+#define mask_or(x, y)     _kor_mask64(x, y)
+#define mask_and(x, y)    _kand_mask64(x, y)
+#define mask_andnot(x, y) _kandn_mask64(y, x)
+#define mask_convert(x)   _mm512_movm_epi8(x)
 
+#define SIMD_LANES 64
+#define SIMD_MASK  0x3F
+
+void skippable_avx512(
+    const int8_t*  restrict strands,
+    const uint8_t* restrict types,
+    const uint8_t* restrict frames,
+    const int min,
+    const int i,
+          uint8_t* restrict skip
+) {
+    skippable_simd(strands, types, frames, min, i, skip);
+}
+
+#endif
 #endif

src/pyrodigal/impl/avx512.pxd

@@ -1,4 +1,7 @@
-from libc.stdint cimport int8_t, uint8_t
+# coding: utf-8
+# cython: language_level=3, linetrace=True, binding=True
 
-cdef extern from "impl/avx512.h" nogil:
-    void skippable_avx512(const int8_t*, const uint8_t*, const uint8_t*, const int, const int, uint8_t*) noexcept
+from ..lib cimport BaseConnectionScorer
+
+cdef class AVX512ConnectionScorer(BaseConnectionScorer):
+    pass

src/pyrodigal/impl/avx512.pyx

@@ -0,0 +1,14 @@
+# coding: utf-8
+# cython: language_level=3, linetrace=True, binding=True
+
+from libc.stdint cimport int8_t, uint8_t
+
+from ..lib cimport BaseConnectionScorer
+
+cdef extern from "avx512.h" nogil:
+    void skippable_avx512(const int8_t*, const uint8_t*, const uint8_t*, const int, const int, uint8_t*) noexcept
+
+cdef class AVX512ConnectionScorer(BaseConnectionScorer):
+    def __cinit__(self):
+        self.skippable = skippable_avx512
+        self.enabled = True

src/scripts/cmake/CythonExtension.cmake

@@ -26,7 +26,7 @@ set(CYTHON_DIRECTIVES
     -E AVX2_BUILD_SUPPORT=$<IF:$<BOOL:${HAVE_AVX2}>,True,False>
     -E NEON_BUILD_SUPPORT=$<IF:$<BOOL:${HAVE_NEON}>,True,False>
     -E MMX_BUILD_SUPPORT=False
-    -E AVX512_BUILD_SUPPORT=False
+    -E AVX512_BUILD_SUPPORT=$<IF:$<BOOL:${HAVE_AVX512}>,True,False>
     -E SYS_IMPLEMENTATION_NAME=$<LOWER_CASE:${Python_INTERPRETER_ID}>
     -E SYS_VERSION_INFO_MAJOR=${Python_VERSION_MAJOR}
     -E SYS_VERSION_INFO_MINOR=${Python_VERSION_MINOR}

src/scripts/cmake/FindAVX512.cmake

@@ -0,0 +1,102 @@
+#.rst:
+# FindAVX512
+# ----------
+#
+# Finds AVX512 support
+#
+# This module can be used to detect AVX512 support in a C compiler.  If
+# the compiler supports AVX512, the flags required to compile with
+# AVX512 support are returned in variables for the different languages.
+# The variables may be empty if the compiler does not need a special
+# flag to support AVX512.
+#
+# The following variables are set:
+#
+# ::
+#
+#    AVX512_C_FLAGS - flags to add to the C compiler for AVX512 support
+#    AVX512_FOUND - true if AVX512 is detected
+#
+#=============================================================================
+
+set(_AVX512_REQUIRED_VARS)
+set(CMAKE_REQUIRED_QUIET_SAVE ${CMAKE_REQUIRED_QUIET})
+set(CMAKE_REQUIRED_QUIET ${AVX512_FIND_QUIETLY})
+
+# sample AVX512 source code to test
+set(AVX512_C_TEST_SOURCE
+"
+#include <immintrin.h>
+void parasail_memset___m256i(__m256i *b, __m256i c, size_t len)
+{
+    size_t i;
+    for (i=0; i<len; ++i) {
+        _mm256_store_si256(&b[i], c);
+    }
+}
+int foo() {
+    __m512i vOne     = _mm512_set1_epi8(1);
+    __mmask64 result = _mm512_cmpeq_epi8_mask(vOne,vOne);
+    return result;
+}
+int main(void) { return (int)foo(); }
+")
+
+# if these are set then do not try to find them again,
+# by avoiding any try_compiles for the flags
+if((DEFINED AVX512_C_FLAGS) OR (DEFINED HAVE_AVX512))
+else()
+  if(WIN32)
+    # MSVC can compile AVX intrinsics without the arch flag, however it
+    # will detect that AVX code is found and "consider using /arch:AVX".
+    set(AVX512_C_FLAG_CANDIDATES
+      "/arch:AVX"
+      "/arch:AVX512")
+  else()
+    set(AVX512_C_FLAG_CANDIDATES
+      #Empty, if compiler automatically accepts AVX512
+      " "
+      #clang, gcc
+      "-mavx512f -mavx512bw"
+    )
+  endif()
+
+  include(CheckCSourceCompiles)
+
+  foreach(FLAG IN LISTS AVX512_C_FLAG_CANDIDATES)
+    set(SAFE_CMAKE_REQUIRED_FLAGS "${CMAKE_REQUIRED_FLAGS}")
+    set(CMAKE_REQUIRED_FLAGS "${FLAG}")
+    unset(HAVE_AVX512 CACHE)
+    if(NOT CMAKE_REQUIRED_QUIET)
+      message(STATUS "Try AVX512 C flag = [${FLAG}]")
+    endif()
+    check_c_source_compiles("${AVX512_C_TEST_SOURCE}" HAVE_AVX512)
+    set(CMAKE_REQUIRED_FLAGS "${SAFE_CMAKE_REQUIRED_FLAGS}")
+    if(HAVE_AVX512)
+      set(AVX512_C_FLAGS_INTERNAL "${FLAG}")
+      break()
+    endif()
+  endforeach()
+
+  unset(AVX512_C_FLAG_CANDIDATES)
+
+  set(AVX512_C_FLAGS "${AVX512_C_FLAGS_INTERNAL}"
+    CACHE STRING "C compiler flags for AVX512 intrinsics")
+endif()
+
+list(APPEND _AVX512_REQUIRED_VARS AVX512_C_FLAGS)
+
+set(CMAKE_REQUIRED_QUIET ${CMAKE_REQUIRED_QUIET_SAVE})
+
+if(_AVX512_REQUIRED_VARS)
+  include(FindPackageHandleStandardArgs)
+
+  find_package_handle_standard_args(AVX512
+                                    REQUIRED_VARS ${_AVX512_REQUIRED_VARS})
+
+  mark_as_advanced(${_AVX512_REQUIRED_VARS})
+
+  unset(_AVX512_REQUIRED_VARS)
+else()
+  message(SEND_ERROR "FindAVX512 requires C or CXX language to be enabled")
+endif()