add linearized b+tree in lsmt

Signed-off-by: Lanzheng Liu <[email protected]>

Author: liulanzheng

.github/workflows/cmake.yml

@@ -51,6 +51,7 @@ jobs:
       working-directory: ${{github.workspace}}/build
       shell: bash
       run: |
+        lscpu
         sudo make install
         sudo cp ${{github.workspace}}/src/example_config/overlaybd-registryv2.json /etc/overlaybd/overlaybd.json
         sudo systemctl enable /opt/overlaybd/overlaybd-tcmu.service
@@ -87,6 +88,7 @@ jobs:
         ls obd_mp
         sudo diff -r --exclude "lost+found" test_data obd_mp
         sudo umount obd_mp
+        cat /var/log/overlaybd.log
     - name: E2E Test turboOCIv1
       working-directory: ${{github.workspace}}/build
       shell: bash

src/overlaybd/lsmt/CMakeLists.txt

@@ -5,6 +5,13 @@ target_include_directories(lsmt_lib PUBLIC
     ${PHOTON_INCLUDE_DIR}
 )
 
+# Check for AVX512F support and add appropriate compiler flags
+include(CheckCXXCompilerFlag)
+check_cxx_compiler_flag("-mavx512f" COMPILER_SUPPORTS_AVX512F)
+if(COMPILER_SUPPORTS_AVX512F)
+    target_compile_options(lsmt_lib PRIVATE "-mavx512f")
+endif()
+
 if(BUILD_TESTING)
   add_subdirectory(test)
 endif()

src/overlaybd/lsmt/index.cpp

@@ -22,6 +22,9 @@
 #include <photon/common/alog.h>
 #include <photon/fs/filesystem.h>
 #include <photon/common/utility.h>
+#ifdef __AVX512F__
+#include <immintrin.h>
+#endif
 using namespace std;
 
 namespace LSMT {
@@ -54,6 +57,123 @@ static inline size_t copy_n(IT begin, IT end, uint64_t end_offset, SegmentMappin
 
 static bool verify_mapping_order(const SegmentMapping *pmappings, size_t n);
 
+bool is_avx512f_supported() {
+    __builtin_cpu_init();
+    return __builtin_cpu_supports("avx512f")  &&
+           __builtin_cpu_supports("avx512dq") &&
+           __builtin_cpu_supports("avx512vl") &&
+           __builtin_cpu_supports("vpclmulqdq");
+}
+
+const static uint32_t ORDER = 8;
+const static uint32_t MAX_LEVEL = 10;
+static constexpr uint32_t NODES_PER_LEVEL[MAX_LEVEL] = {8, 72, 648, 5832, 52488, 472392, 4251528, 38263752, 344373768, 3099363912};
+static constexpr uint32_t LEVEL_START_ID[MAX_LEVEL] =  {0,  8,  80,  728,  6560,  59048,  531440,  4782968,  43046720,  387420488};
+
+class LinearizedBptree: public Object {
+public:
+    uint64_t N;
+    uint64_t *node = nullptr;
+    int32_t DEPTH = -1;
+
+    LinearizedBptree() {
+    }
+
+    ~LinearizedBptree() {
+        free(node);
+    }
+
+    virtual int build(const SegmentMapping* pbegin, const SegmentMapping* pend) {
+        if (pbegin == pend) {
+            LOG_ERROR_RETURN(EINVAL, -1, "linearized bptree not used: empty mapping");
+        }
+        if (pbegin->offset != 0) {
+            // In a real file system, mapping offset starts from 0. skip for some ut.
+            LOG_ERROR_RETURN(EINVAL, -1, "linearized bptree not used: invalid start offset");
+        }
+        size_t mapping_size = pend - pbegin;
+        for (uint32_t i = 0; i < MAX_LEVEL; i++)
+            if (NODES_PER_LEVEL[i] >= mapping_size) {
+                DEPTH = i+1;
+                break;
+            }
+        if (DEPTH == -1) {
+            LOG_ERROR_RETURN(EINVAL, -1, "linearized bptree not used: too many mappings");
+        }
+
+        N = (LEVEL_START_ID[DEPTH-1] + mapping_size + ORDER - 1) / ORDER * ORDER;
+        LOG_INFO("building Linearized B+tree ", VALUE(DEPTH), VALUE(mapping_size), VALUE(N));
+        auto ret = posix_memalign((void**)&node, 64, N*sizeof(uint64_t));
+        if (ret != 0) {
+            LOG_ERRNO_RETURN(ENOBUFS, -1, "linearized bptree not used: failed to alloc memory");
+        }
+
+        uint32_t leaf_start = LEVEL_START_ID[DEPTH - 1];
+        uint32_t leaf_size = NODES_PER_LEVEL[DEPTH - 1];
+
+        uint32_t p = leaf_start;
+
+        for (auto mp = pbegin; mp < pend; mp++, p++) {
+            node[p] = mp->offset;
+        }
+        for (; p < N; p++)
+            node[p] = -1;
+
+        auto G = ORDER;
+        for (auto level = DEPTH-1; level > 0; level--) {
+            auto pos = LEVEL_START_ID[level - 1];
+            for (uint32_t i = 0; i < leaf_size; i += G * (ORDER + 1)) {
+                for (uint32_t j = 1; j <= ORDER; j++) {
+                    uint32_t lower_id = leaf_start + i + G * j;
+                    node[pos++] = (lower_id < N) ? node[lower_id] : -1;
+                }
+            }
+            G *= (ORDER + 1);
+        }
+        LOG_INFO("building Linearized B+tree done");
+        return 0;
+    }
+
+    virtual uint32_t branchfree_inner_search(const uint64_t *base, uint64_t x) const {
+        return std::upper_bound(base, base + 8, x) - base;
+    }
+
+    __attribute__((always_inline)) uint32_t search(const uint64_t x) const {
+        uint32_t res = 0;
+#pragma GCC unroll 20
+        for (int i = DEPTH; i > 1; --i) {
+            uint32_t c = branchfree_inner_search(node + res, x);
+            res = (ORDER+1)*res + (c+1)*ORDER;
+        }
+        res += branchfree_inner_search(node + res, x);
+        res = res - 1 - LEVEL_START_ID[DEPTH-1];
+        return res;
+    }
+};
+
+class LinearizedBptreeAcc: public LinearizedBptree {
+public:
+    int build(const SegmentMapping* pbegin, const SegmentMapping* pend) override {
+        auto res = LinearizedBptree::build(pbegin, pend);
+        if (res < 0)
+            return res;
+        LOG_INFO("using Accelerated Linearized B+tree");
+        return 0;
+    }
+    uint32_t branchfree_inner_search(const uint64_t *base, uint64_t x) const override {
+#ifdef __AVX512F__
+        __m512i vx = _mm512_set1_epi64(x);
+        __m512i data = _mm512_load_si512(base);
+        uint8_t mask = _mm512_cmp_epu64_mask(vx, data, _MM_CMPINT_GE);
+        return __builtin_popcount(mask);
+#else
+        return std::upper_bound(base, base + 8, x) - base;
+#endif
+
+    }
+};
+
+
 class Index : public IMemoryIndex {
 public:
     bool ownership = false;
@@ -62,6 +182,7 @@ class Index : public IMemoryIndex {
     const SegmentMapping *pend = nullptr;
     uint64_t alloc_blk = 0;
     uint64_t virtual_size = 0;
+    LinearizedBptree *lbt = nullptr;
 
     inline void get_alloc_blks() {
         for (auto m : mapping) {
@@ -72,6 +193,7 @@ class Index : public IMemoryIndex {
         if (ownership) {
             delete[] pbegin;
         }
+        safe_delete(lbt);
     }
     Index(const SegmentMapping *pmappings = nullptr, size_t n = 0, bool ownership = true,
           uint64_t vsize = 0)
@@ -120,7 +242,15 @@ class Index : public IMemoryIndex {
     virtual size_t lookup(Segment s, /* OUT */ SegmentMapping *pm, size_t n) const override {
         if (s.length == 0)
             return 0;
-        auto lb = std::lower_bound(pbegin, pend, s);
+        const SegmentMapping *lb;
+        if (lbt) {
+            lb = pbegin + lbt->search(s.offset);
+            if (lb->end() <= s.offset)
+                lb++;
+        } else {
+            lb = std::lower_bound(pbegin, pend, s);
+        }
+
         auto m = copy_n(lb, pend, s.end(), pm, n);
         trim_edge_mappings(pm, m, s);
         return m;
@@ -155,6 +285,16 @@ class Index : public IMemoryIndex {
     }
 
     UNIMPLEMENTED_POINTER(IMemoryIndex  *make_read_only_index() const override);
+
+    int build_lineriazed_bptree() {
+        LinearizedBptree *tree = is_avx512f_supported() ? new LinearizedBptreeAcc() : new LinearizedBptree();
+        if (tree->build(pbegin, pend) < 0) {
+            delete(tree);
+            return -1;
+        }
+        lbt = tree;
+        return 0;
+    }
 };
 
 class LevelIndex : public Index {
@@ -694,6 +834,8 @@ IMemoryIndex *merge_memory_indexes(const IMemoryIndex **pindexes, size_t n) {
     auto pi = (const Index **)pindexes;
     mapping.reserve(pi[0]->size());
     merge_indexes(0, mapping, pi, n, 0, UINT64_MAX);
-    return new Index(std::move(mapping), pindexes[0]->vsize());
+    auto idx = new Index(std::move(mapping), pindexes[0]->vsize());
+    idx->build_lineriazed_bptree();
+    return idx;
 }
 } // namespace LSMT

src/overlaybd/lsmt/index.h

@@ -29,9 +29,9 @@ IMemoryIndex -> IMemoryIndex0 -> IComboIndex -> Index0 ( set<SegmentMap> ) -> Co
 #include <sys/types.h>
 
 namespace LSMT {
-struct Segment {          // 48 + 18 == 64
-    uint64_t offset : 50; // offset (0.5 PB if in sector)
-    uint32_t length : 14; // length (8MB if in sector)
+struct Segment {          
+    uint64_t offset : 50;
+    uint32_t length : 14;
     const static uint64_t MAX_OFFSET = (1UL << 50) - 1;
     const static uint32_t MAX_LENGTH = (1 << 14) - 1;
     const static uint64_t INVALID_OFFSET = MAX_OFFSET;

src/overlaybd/lsmt/test/CMakeLists.txt

@@ -13,3 +13,15 @@ add_test(
   COMMAND ${EXECUTABLE_OUTPUT_PATH}/lsmt_test
 )
 
+include(CheckCXXCompilerFlag)
+check_cxx_compiler_flag("-mavx512f" COMPILER_SUPPORTS_AVX512F)
+if (COMPILER_SUPPORTS_AVX512F)
+  add_executable(lsmt_test_lbt test.cpp)
+  target_include_directories(lsmt_test_lbt PUBLIC ${PHOTON_INCLUDE_DIR})
+  target_link_libraries(lsmt_test_lbt gtest gtest_main gflags pthread photon_static overlaybd_lib)
+  target_compile_options(lsmt_test_lbt PRIVATE "-mavx512f")
+  add_test(
+    NAME lsmt_test_lbt
+    COMMAND ${EXECUTABLE_OUTPUT_PATH}/lsmt_test_lbt
+  )
+endif()