check loongarch lasx and enable (#4820)

Tencent · Jun 27, 2023 · 6c21b08 · 6c21b08
1 parent 43aba6b
commit 6c21b08
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 2 deletions.
diff --git a/.ci/test-coverage.yml b/.ci/test-coverage.yml
@@ -891,7 +891,7 @@ jobs:
       run: |
         export LOONGARCH64_ROOT_PATH=${{ci.workspace}}/cross-tools
         mkdir build && cd build
-        cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/loongarch64-unknown-linux-gnu.toolchain.cmake -DCMAKE_BUILD_TYPE=debug -DNCNN_COVERAGE=ON -DNCNN_RUNTIME_CPU=OFF -DNCNN_LSX=ON -DNCNN_OPENMP=${{matrix.OPENMP}} -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF -DNCNN_BUILD_TESTS=ON ..
+        cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/loongarch64-unknown-linux-gnu.toolchain.cmake -DCMAKE_BUILD_TYPE=debug -DNCNN_COVERAGE=ON -DNCNN_RUNTIME_CPU=OFF -DNCNN_LSX=ON -DNCNN_LASX=OFF -DNCNN_OPENMP=${{matrix.OPENMP}} -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF -DNCNN_BUILD_TESTS=ON ..
         cmake --build . -j $(nproc)
     - name: test
       run: |

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -303,10 +303,21 @@ elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(mips)")
 elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(loongarch64|loongarch32)")
     set(NCNN_TARGET_ARCH loongarch)
 
-    check_cxx_compiler_flag("-mlsx" NCNN_COMPILER_SUPPORT_LOONGARCH_LSX)
+    set(CMAKE_REQUIRED_FLAGS "-mlsx")
+    check_cxx_source_compiles("#include <lsxintrin.h>\nint main() { __m128 _s, _a, _b, _c; _s = __lsx_vfmadd_s(_a, _b, _c); return 0; }" NCNN_COMPILER_SUPPORT_LOONGARCH_LSX)
+
+    set(CMAKE_REQUIRED_FLAGS "-mlasx")
+    check_cxx_source_compiles("#include <lasxintrin.h>\nint main() { __m256 _s, _a, _b, _c; _s = __lasx_xvfmadd_s(_a, _b, _c); return 0; }" NCNN_COMPILER_SUPPORT_LOONGARCH_LASX)
+
+    unset(CMAKE_REQUIRED_FLAGS)
 
     if(NCNN_COMPILER_SUPPORT_LOONGARCH_LSX)
         option(NCNN_LSX "optimize loongarch platform with lsx extension" ON)
+        if(NCNN_COMPILER_SUPPORT_LOONGARCH_LASX)
+            option(NCNN_LASX "optimize loongarch platform with lasx extension" ON)
+        else()
+            message(WARNING "The compiler does not support lasx extension. NCNN_LASX will be OFF.")
+        endif()
     else()
         message(WARNING "The compiler does not support lsx extension. NCNN_LSX will be OFF.")
     endif()

diff --git a/cmake/ncnn_add_layer.cmake b/cmake/ncnn_add_layer.cmake
@@ -273,6 +273,9 @@ macro(ncnn_add_layer class)
     endif()
 
     if(NCNN_TARGET_ARCH STREQUAL "loongarch")
+        if(NCNN_RUNTIME_CPU AND NCNN_LASX)
+            ncnn_add_arch_opt_layer(${class} lasx "-mlasx -mlsx")
+        endif()
         if(NCNN_RUNTIME_CPU AND NCNN_LSX)
             ncnn_add_arch_opt_layer(${class} lsx "-mlsx")
         endif()

diff --git a/cmake/ncnn_generate_lasx_source.cmake b/cmake/ncnn_generate_lasx_source.cmake
@@ -0,0 +1,14 @@
+
+# must define SRC DST CLASS
+
+file(READ ${SRC} source_data)
+
+# replace
+string(TOUPPER ${CLASS} CLASS_UPPER)
+string(TOLOWER ${CLASS} CLASS_LOWER)
+
+string(REGEX REPLACE "LAYER_${CLASS_UPPER}_LOONGARCH_H" "LAYER_${CLASS_UPPER}_LOONGARCH_LASX_H" source_data "${source_data}")
+string(REGEX REPLACE "${CLASS}_loongarch" "${CLASS}_loongarch_lasx" source_data "${source_data}")
+string(REGEX REPLACE "#include \"${CLASS_LOWER}_loongarch.h\"" "#include \"${CLASS_LOWER}_loongarch_lasx.h\"" source_data "${source_data}")
+
+file(WRITE ${DST} "${source_data}")
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -484,6 +484,9 @@ endif()
 if(NCNN_TARGET_ARCH STREQUAL "loongarch")
     if(NOT NCNN_RUNTIME_CPU AND NCNN_LSX)
         target_compile_options(ncnn PRIVATE -mlsx)
+        if(NCNN_LASX)
+            target_compile_options(ncnn PRIVATE -mlasx)
+        endif()
     endif()
 endif()
 

diff --git a/src/layer.cpp b/src/layer.cpp
@@ -253,6 +253,13 @@ Layer* create_layer(int index)
     }
     else
 #endif // NCNN_RUNTIME_CPU && NCNN_AVX
+#if NCNN_RUNTIME_CPU && NCNN_LASX
+    if (ncnn::cpu_support_loongarch_lasx())
+    {
+        layer_creator = layer_registry_lasx[index].creator;
+    }
+    else
+#endif // NCNN_RUNTIME_CPU && NCNN_LASX
 #if NCNN_RUNTIME_CPU && NCNN_LSX
     if (ncnn::cpu_support_loongarch_lsx())
     {

diff --git a/src/layer_registry.h.in b/src/layer_registry.h.in
@@ -28,6 +28,12 @@ static const layer_registry_entry layer_registry_msa[] = {
 };
 #endif // NCNN_RUNTIME_CPU && NCNN_MSA
 
+#if NCNN_RUNTIME_CPU && NCNN_LASX
+static const layer_registry_entry layer_registry_lasx[] = {
+@layer_registry_lasx@
+};
+#endif // NCNN_RUNTIME_CPU && NCNN_LASX
+
 #if NCNN_RUNTIME_CPU && NCNN_LSX
 static const layer_registry_entry layer_registry_lsx[] = {
 @layer_registry_lsx@