feat: support neon, sse simd and dynamic dispatch (bytedance#56)

# Main changes - support both static dispatch and dynamic dispatch - support neon and sse architecture
moyin1004 · Apr 11, 2023 · 50514ec · 50514ec
1 parent e3038ea
commit 50514ec
Show file tree

Hide file tree

Showing 60 changed files with 3,955 additions and 1,193 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -7,13 +7,12 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        llvm_version:
-          - '11'
-          - '13.0.0'
-          - '15'
-        tool:
-          - 'cmake'
-          - 'bazel'
+        llvm_version: ['11', '13.0.0', '15']
+        tool: ['cmake', 'bazel']
+        arch: [westmere, haswell]
+        exclude:
+            - tool: 'cmake'
+              arch: westmere
     env:
       CC: clang
       CXX: clang++
@@ -38,10 +37,10 @@ jobs:
         with:
           cmake-version: "latest"
 
-      - name: Run Test Use Bazel
+      - name: Run ${{ matrix.arch }} Test Use Bazel
         if: matrix.tool == 'bazel'
         run : |
-          bash ./scripts/unittest.sh -c
+          bash ./scripts/unittest.sh -c --arch=${{ matrix.arch }}
 
       - name: Run Test Use CMake
         if: matrix.tool == 'cmake'
@@ -53,12 +52,19 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        gcc_version:
-          - '9'
-          - '12'
-        tool:
-          - 'cmake'
-          - 'bazel'
+        gcc_version: ['9', '12']
+        tool: ['cmake', 'bazel']
+        arch: [westmere, haswell]
+        dispatch: [static, dynamic]
+        exclude:
+          - tool: 'bazel'
+            arch: haswell
+            dispatch: dynamic
+          - tool: 'cmake'
+            dispatch: dynamic
+          - tool: 'cmake'
+            arch: westmere
+
     env:
       CC: gcc-${{ matrix.gcc_version }}
       CXX: g++-${{ matrix.gcc_version}}
@@ -82,10 +88,10 @@ jobs:
         with:
           cmake-version: "latest"
 
-      - name: Run Test Use Bazel
+      - name: Run ${{ matrix.arch }} ${{ matrix.dispatch }} Test Use Bazel
         if: matrix.tool == 'bazel'
         run : |
-          bash ./scripts/unittest.sh -g
+          bash ./scripts/unittest.sh -g --arch=${{ matrix.arch }} --dispatch=${{ matrix.dispatch }}
 
       - name: Run Test Use CMake
         if: matrix.tool == 'cmake'

diff --git a/BUILD.bazel b/BUILD.bazel
@@ -1,8 +1,73 @@
 package(default_visibility = ["//visibility:public"])
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
 
 common_copts = ['-mavx2', '-mbmi', '-mpclmul']
 sanitize_copts = ['-fsanitize=address,undefined', '-fsanitize-recover=address']
 
+avx2_copts = ['-mavx2', '-mbmi', '-mpclmul']
+sse_copts = ['-msse', '-msse2', '-msse4.1', '-mpclmul']
+arm_copts = ['-march-armv8-a']
+benchmark_copts = ['-O3', '-DNDEBUG', '-std=c++17']
+static_dispatch_copts = []
+dynamic_dispatch_copts = ['-DSONIC_DYNAMIC_DISPATCH=1']
+no_san_linkopts = []
+gcc_san_linkopts = ['-lasan']
+clang_san_linkopts = ['-fsanitize-link-c++-runtime']
+
+string_flag(
+    name = "sonic_arch",
+    build_setting_default = "haswell",
+)
+
+string_flag(
+    name = "sonic_dispatch",
+    build_setting_default = "static",
+)
+
+string_flag(
+    name = "sonic_sanitizer",
+    build_setting_default = "no",
+)
+
+config_setting(
+    name = "arm_build",
+    flag_values = {":sonic_arch": "arm"},
+)
+
+config_setting(
+    name = "sse_build",
+    flag_values = {":sonic_arch": "westmere"},
+)
+config_setting(
+    name = "avx2_build",
+    flag_values = {":sonic_arch": "haswell"},
+)
+
+config_setting(
+    name = "static_dispatch",
+    flag_values = {":sonic_dispatch": "static"},
+)
+
+config_setting(
+    name = "dynamic_dispatch",
+    flag_values = {":sonic_dispatch": "dynamic"},
+)
+
+config_setting(
+    name = "no_san",
+    flag_values = {":sonic_sanitizer": "no"},
+)
+
+config_setting(
+    name = "gcc_san",
+    flag_values = {":sonic_sanitizer": "gcc"},
+)
+
+config_setting(
+    name = "clang_san",
+    flag_values = {":sonic_sanitizer": "clang"},
+)
+
 cc_library(
     name = "string_view",
     hdrs = glob(["include/thirdparty/**/*.h"]),
@@ -35,12 +100,21 @@ cc_binary(
         "@simdjson",
         "@jsoncpp//:jsoncpp",
     ],
-    copts = common_copts + ['-DNDEBUG', '-std=c++17'],
+    # copts = common_copts + ['-DNDEBUG', '-std=c++17'],
+    copts = select({
+                "arm_build": arm_copts + benchmark_copts,
+                "sse_build": sse_copts + benchmark_copts,
+                "avx2_build": avx2_copts + benchmark_copts,
+            }) +\
+            select({
+                "static_dispatch": static_dispatch_copts,
+                "dynamic_dispatch": dynamic_dispatch_copts,
+            }),
     linkopts = ['-lstdc++fs'],
 )
 
 cc_test(
-    name = "unittest-gcc",
+    name = "unittest",
     srcs = glob([
         "tests/*.cpp",
         "include/sonic/*",
@@ -51,10 +125,49 @@ cc_test(
         "@gtest//:gtest_main",
     ],
     data = glob([ "testdata/*.json"]),
+    linkopts = [ '-lstdc++fs', '-fstack-protector-all',] +\
+        select({
+            "no_san": [],
+            "gcc_san": sanitize_copts + gcc_san_linkopts,
+            "clang_san": sanitize_copts + clang_san_linkopts,
+        }),
+    copts = [
+        '-O3', '-g', '-UNDEBUG', '-std=c++14',
+        '-fstack-protector-all',
+        '-Iinclude', '-Wall', '-Wextra', '-Werror',
+    ]  + \
+        select({
+            "no_san": [],
+            "gcc_san": sanitize_copts,
+            "clang_san": sanitize_copts,
+        }) + \
+        select({
+            "arm_build": arm_copts + benchmark_copts,
+            "sse_build": sse_copts + benchmark_copts,
+            "avx2_build": avx2_copts + benchmark_copts,
+        }) +\
+        select({
+            "static_dispatch": static_dispatch_copts,
+            "dynamic_dispatch": dynamic_dispatch_copts,
+        }),
+)
+
+cc_test(
+    name = "unittest-clang",
+    srcs = glob([
+            "tests/*.h",
+            "tests/*.cpp",
+            "include/sonic/*",
+            "include/sonic/**/*",
+        ]),
+    deps = [
+        ":string_view",
+        "@gtest//:gtest_main",
+    ],
+    data = glob([ "testdata/*.json"]),
     linkopts = sanitize_copts + [
         '-lstdc++fs',
         '-fstack-protector-all',
-        '-lasan'
     ],
     copts = common_copts + sanitize_copts + [
         '-O3', '-g', '-UNDEBUG', '-std=c++14',
@@ -64,14 +177,14 @@ cc_test(
 )
 
 cc_test(
-    name = "unittest-clang",
+    name = "unittest-arm",
     srcs = glob([
             "tests/*.h",
             "tests/*.cpp",
             "include/sonic/*",
             "include/sonic/**/*",
         ]),
-    deps = [
+    deps = [ 
         ":string_view",
         "@gtest//:gtest_main",
     ],
@@ -81,8 +194,33 @@ cc_test(
         '-fstack-protector-all',
         '-fsanitize-link-c++-runtime'
     ],
-    copts = common_copts + sanitize_copts + [
-        '-O3', '-g', '-UNDEBUG', '-std=c++14',
+    copts = sanitize_copts + [
+        '-O3', '-g', '-UNDEBUG', '-std=c++14', '-march=armv8-a',
+        '-fstack-protector-all',
+        '-Iinclude', '-Wall', '-Wextra', '-Werror',
+    ],
+)
+
+cc_test(
+    name = "unittest-sse",
+    srcs = glob([
+            "tests/*.h",
+            "tests/*.cpp",
+            "include/sonic/*",
+            "include/sonic/**/*",
+        ]),
+    deps = [ 
+        ":string_view",
+        "@gtest//:gtest_main",
+    ],
+    data = glob([ "testdata/*.json"]),
+    linkopts = sanitize_copts + [
+        '-lstdc++fs',
+        '-fstack-protector-all',
+        '-fsanitize-link-c++-runtime'
+    ],
+    copts = sanitize_copts + [
+        '-O3', '-g', '-UNDEBUG', '-std=c++14', '-march=westmere',
         '-fstack-protector-all',
         '-Iinclude', '-Wall', '-Wextra', '-Werror',
     ],
@@ -112,3 +250,4 @@ cc_test(
         '-fprofile-arcs', '-ftest-coverage',
     ],
 )
+
diff --git a/WORKSPACE b/WORKSPACE
@@ -1,5 +1,19 @@
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "new_git_repository")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "bazel_skylib",
+    sha256 = "b8a1527901774180afc798aeb28c4634bdccf19c4d98e7bdd1ce79d1fe9aaad7",
+    urls = [
+        "https://mirror.bazel.build/github.com/bazelbuild/bazel-skylib/releases/download/1.4.1/bazel-skylib-1.4.1.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.4.1/bazel-skylib-1.4.1.tar.gz",
+    ],
+)
+
+load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace")
+
+bazel_skylib_workspace()
 
 git_repository(
     name = "google_benchmark",

diff --git a/bazel/cJSON.BUILD b/bazel/cJSON.BUILD
@@ -3,5 +3,6 @@ cc_library(
   name = "cJSON",
   srcs = ["cJSON.c"],
   hdrs = ["cJSON.h"],
-  copts = ['-O3' ,'-DNDEBUG', '-march=haswell'],
+  copts = ['-O3' ,'-DNDEBUG',],
 )
+
diff --git a/bazel/yyjson.BUILD b/bazel/yyjson.BUILD
@@ -4,5 +4,6 @@ cc_library(
   srcs = ["src/yyjson.c"],
   hdrs = ["src/yyjson.h"],
   includes = ["src"],
-  copts = ['-O3', '-DNDEBUG', '-march=haswell', '-g'],
+  copts = ['-O3', '-DNDEBUG', '-g'],
 )
+
diff --git a/include/sonic/allocator.h b/include/sonic/allocator.h
@@ -89,7 +89,17 @@ class SpinLock {
         break;
       }
       while (lock_.load(std::memory_order_relaxed)) {
+        // use pause or yield instruction will slow down lock acquisition
+        // on contended locks.
+#ifndef SONIC_SPINLOCK_NO_PAUSE
+
+#if defined(__x86_64__) || defined(_M_AMD64)
         __builtin_ia32_pause();
+#elif defined(__aarch64__) || defined(_M_ARM64)
+        asm volatile("yield");
+#endif
+
+#endif
       }
     }
   }

diff --git a/include/sonic/dom/dynamicnode.h b/include/sonic/dom/dynamicnode.h
@@ -27,8 +27,6 @@
 #include "sonic/dom/type.h"
 #include "sonic/error.h"
 #include "sonic/internal/ftoa.h"
-#include "sonic/internal/itoa.h"
-#include "sonic/internal/quote.h"
 #include "sonic/writebuffer.h"
 
 namespace sonic_json {

diff --git a/include/sonic/dom/handler.h b/include/sonic/dom/handler.h
@@ -19,7 +19,7 @@
 #include <string>
 
 #include "sonic/dom/type.h"
-#include "sonic/internal/haswell.h"
+#include "sonic/internal/arch/simd_base.h"
 #include "sonic/string_view.h"
 #include "sonic/writebuffer.h"
 
@@ -153,7 +153,7 @@ class SAXHandler {
     if (pairs) {
       void *mem = obj.template containerMalloc<MemberType>(pairs, *alloc_);
       obj.setChildren(mem);
-      internal::haswell::xmemcpy<sizeof(MemberType)>(
+      internal::Xmemcpy<sizeof(MemberType)>(
           (void *)obj.getObjChildrenFirstUnsafe(), (void *)(&obj + 1), pairs);
     } else {
       obj.setChildren(nullptr);
@@ -169,7 +169,7 @@ class SAXHandler {
     arr.setLength(count, kArray);
     if (count) {
       arr.setChildren(arr.template containerMalloc<NodeType>(count, *alloc_));
-      internal::haswell::xmemcpy<sizeof(NodeType)>(
+      internal::Xmemcpy<sizeof(NodeType)>(
           (void *)arr.getArrChildrenFirstUnsafe(), (void *)(&arr + 1), count);
     } else {
       arr.setChildren(nullptr);
@@ -239,7 +239,7 @@ class LazySAXHandler {
     arr.setLength(count, kArray);
     if (count) {
       arr.setChildren(arr.template containerMalloc<NodeType>(count, *alloc_));
-      internal::haswell::xmemcpy<sizeof(NodeType)>(
+      internal::Xmemcpy<sizeof(NodeType)>(
           (void *)arr.getArrChildrenFirstUnsafe(), (void *)(&arr + 1), count);
       stack_.Pop<NodeType>(count);
     } else {
@@ -254,7 +254,7 @@ class LazySAXHandler {
     if (pairs) {
       void *mem = obj.template containerMalloc<MemberType>(pairs, *alloc_);
       obj.setChildren(mem);
-      internal::haswell::xmemcpy<sizeof(MemberType)>(
+      internal::Xmemcpy<sizeof(MemberType)>(
           (void *)obj.getObjChildrenFirstUnsafe(), (void *)(&obj + 1), pairs);
       stack_.Pop<MemberType>(pairs);
     } else {