Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

implement of the engine and rewrite factory #209

Merged
merged 1 commit into from
Dec 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions examples/cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ target_link_libraries (example_fresh_hnsw vsag)
add_executable (example_diskann example_diskann.cpp)
target_link_libraries (example_diskann vsag_static)

add_executable (example_engine example_engine.cpp)
target_link_libraries (example_engine vsag_static)

add_executable (simple_hnsw simple_hnsw.cpp)
target_link_libraries (simple_hnsw vsag_static)

Expand Down
134 changes: 134 additions & 0 deletions examples/cpp/example_engine.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@

// Copyright 2024-present the vsag project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <vsag/vsag.h>

#include <iostream>

class SimpleAllocator : public vsag::Allocator {
public:
SimpleAllocator() = default;
~SimpleAllocator() override = default;

public:
std::string
Name() override {
return "DefaultAllocator";
}

void*
Allocate(size_t size) override {
auto ptr = malloc(size);
allocate_count_++;
return ptr;
}

void
Deallocate(void* p) override {
free(p);
}

void*
Reallocate(void* p, size_t size) override {
auto ptr = realloc(p, size);
return ptr;
}

void
Print() {
std::cout << "allocate times : " << allocate_count_ << std::endl;
}
uint64_t allocate_count_{0};
};

int
main(int argc, char** argv) {
vsag::init();

int64_t num_vectors = 1000;
int64_t dim = 128;

// prepare ids and vectors
auto ids = new int64_t[num_vectors];
auto vectors = new float[dim * num_vectors];

std::mt19937 rng(47);
std::uniform_real_distribution<float> distrib_real;
for (int64_t i = 0; i < num_vectors; ++i) {
ids[i] = i;
}
for (int64_t i = 0; i < dim * num_vectors; ++i) {
vectors[i] = distrib_real(rng);
}

// create index
auto hnsw_build_paramesters = R"(
{
"dtype": "float32",
"metric_type": "l2",
"dim": 128,
"hnsw": {
"max_degree": 16,
"ef_construction": 100
}
}
)";

auto* allocator = new SimpleAllocator();

{
vsag::Resource resource(allocator);
vsag::Engine engine(&resource);
auto index = engine.CreateIndex("hnsw", hnsw_build_paramesters).value();
auto base = vsag::Dataset::Make();
base->NumElements(num_vectors)->Dim(dim)->Ids(ids)->Float32Vectors(vectors)->Owner(false);
index->Build(base);

// prepare a query vector
// memory will be released by query the dataset since owner is set to true when creating the query.
auto query_vector = new float[dim];
for (int64_t i = 0; i < dim; ++i) {
query_vector[i] = distrib_real(rng);
}

// search on the index
auto hnsw_search_parameters = R"(
{
"hnsw": {
"ef_search": 100
}
}
)";
int64_t topk = 10;
auto query = vsag::Dataset::Make();
query->NumElements(1)->Dim(dim)->Float32Vectors(query_vector)->Owner(true);
auto result = index->KnnSearch(query, topk, hnsw_search_parameters).value();

// print the results
std::cout << "results: " << std::endl;
for (int64_t i = 0; i < result->GetDim(); ++i) {
std::cout << result->GetIds()[i] << ": " << result->GetDistances()[i] << std::endl;
}
engine.Shutdown();
allocator->Print();
}

// free memory
delete[] ids;
delete[] vectors;
delete allocator;

return 0;
}
7 changes: 5 additions & 2 deletions include/vsag/engine.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,19 @@

#include <memory>

#include "dataset.h"
#include "index.h"
#include "resource.h"

namespace vsag {
class Engine {
public:
explicit Engine(Resource* resource = nullptr);
explicit Engine();

explicit Engine(Resource* resource);

void
Shutdown();
Shutdown(); // like ~Engine(), but will warn whether the resources are still reference outside

tl::expected<std::shared_ptr<Index>, Error>
CreateIndex(const std::string& name, const std::string& parameters);
Expand Down
9 changes: 6 additions & 3 deletions include/vsag/resource.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,15 +20,18 @@
#include "allocator.h"

namespace vsag {
struct Resource {
class Resource {
public:
explicit Resource(Allocator* allocator = nullptr);

virtual ~Resource() = default;

virtual std::shared_ptr<Allocator>
GetAllocator() {
return this->allocator;
}

public:
std::shared_ptr<Allocator> allocator;

// std::shared_ptr<> thread_pool;
};
} // namespace vsag
1 change: 1 addition & 0 deletions include/vsag/vsag.h
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ init();
#include "bitset.h"
#include "constants.h"
#include "dataset.h"
#include "engine.h"
#include "errors.h"
#include "expected.hpp"
#include "factory.h"
Expand Down
12 changes: 6 additions & 6 deletions src/algorithm/hgraph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -53,14 +53,14 @@ next_multiple_of_power_of_two(uint64_t x, uint64_t n) {
HGraph::HGraph(const JsonType& index_param, const vsag::IndexCommonParam& common_param) noexcept
: index_param_(index_param),
common_param_(common_param),
label_lookup_(common_param.allocator_),
label_op_mutex_(MAX_LABEL_OPERATION_LOCKS, common_param_.allocator_),
neighbors_mutex_(0, common_param_.allocator_),
route_graphs_(common_param.allocator_),
labels_(common_param.allocator_) {
label_lookup_(common_param.allocator_.get()),
label_op_mutex_(MAX_LABEL_OPERATION_LOCKS, common_param.allocator_.get()),
neighbors_mutex_(0, common_param.allocator_.get()),
route_graphs_(common_param.allocator_.get()),
labels_(common_param.allocator_.get()) {
this->dim_ = common_param.dim_;
this->metric_ = common_param.metric_;
this->allocator_ = common_param.allocator_;
this->allocator_ = common_param.allocator_.get();
jiaweizone marked this conversation as resolved.
Show resolved Hide resolved
}

tl::expected<void, Error>
Expand Down
2 changes: 1 addition & 1 deletion src/data_cell/flatten_datacell.h
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ template <typename QuantTmpl, typename IOTmpl>
FlattenDataCell<QuantTmpl, IOTmpl>::FlattenDataCell(const JsonType& quantization_param,
const JsonType& io_param,
const IndexCommonParam& common_param)
: allocator_(common_param.allocator_) {
: allocator_(common_param.allocator_.get()) {
this->quantizer_ = std::make_shared<QuantTmpl>(quantization_param, common_param);
this->io_ = std::make_shared<IOTmpl>(io_param, common_param);
this->code_size_ = quantizer_->GetCodeSize();
Expand Down
8 changes: 5 additions & 3 deletions src/data_cell/flatten_datacell_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,15 @@
#include "flatten_datacell.h"

#include <algorithm>
#include <utility>

#include "catch2/catch_template_test_macros.hpp"
#include "default_allocator.h"
#include "fixtures.h"
#include "flatten_interface_test.h"
#include "io/io_headers.h"
#include "quantization/quantizer_headers.h"
#include "safe_allocator.h"

using namespace vsag;

Expand All @@ -36,7 +38,7 @@ TestFlattenDataCell(int dim,
auto counts = {100, 1000};
IndexCommonParam common;
common.dim_ = dim;
common.allocator_ = allocator.get();
common.allocator_ = std::move(allocator);
common.metric_ = metric;
for (auto count : counts) {
auto flatten =
Expand Down Expand Up @@ -67,7 +69,7 @@ TestFlattenDataCellFP32(int dim,
}

TEST_CASE("fp32", "[ut][flatten_data_cell]") {
auto allocator = std::make_shared<DefaultAllocator>();
auto allocator = SafeAllocator::FactoryDefaultAllocator();
auto fp32_param = JsonType::parse("{}");
auto io_param = JsonType::parse("{}");
auto dims = {8, 64, 512};
Expand Down Expand Up @@ -96,7 +98,7 @@ TestFlattenDataCellSQ8(int dim,
}

TEST_CASE("sq8", "[ut][flatten_data_cell]") {
auto allocator = std::make_shared<DefaultAllocator>();
auto allocator = SafeAllocator::FactoryDefaultAllocator();
auto sq8_param = JsonType::parse("{}");
auto io_param = JsonType::parse("{}");
auto dims = {32, 64, 512};
Expand Down
5 changes: 3 additions & 2 deletions src/data_cell/graph_datacell_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include "fmt/format-inl.h"
#include "graph_interface_test.h"
#include "io/io_headers.h"
#include "safe_allocator.h"
using namespace vsag;

template <typename IOTemp>
Expand All @@ -37,7 +38,7 @@ TestGraphDataCell(const JsonType& graph_param,
}

TEST_CASE("graph basic test", "[ut][graph_datacell]") {
auto allocator = std::make_shared<DefaultAllocator>();
auto allocator = SafeAllocator::FactoryDefaultAllocator();
auto dims = {32, 64};
auto max_degrees = {5, 12, 24, 32, 64, 128};
auto max_capacities = {1, 100, 10000, 10'000'000, 32'179'837};
Expand All @@ -58,7 +59,7 @@ TEST_CASE("graph basic test", "[ut][graph_datacell]") {
for (auto dim : dims) {
IndexCommonParam param;
param.dim_ = dim;
param.allocator_ = allocator.get();
param.allocator_ = allocator;
for (auto& gp : graph_params) {
TestGraphDataCell<MemoryIO>(gp, io_param, param);
TestGraphDataCell<MemoryBlockIO>(gp, io_param, param);
Expand Down
3 changes: 2 additions & 1 deletion src/data_cell/graph_interface_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,12 +20,13 @@
#include "catch2/catch_test_macros.hpp"
#include "default_allocator.h"
#include "fixtures.h"
#include "safe_allocator.h"

using namespace vsag;

void
GraphInterfaceTest::BasicTest(uint64_t max_id, uint64_t count, GraphInterfacePtr other) {
auto allocator = std::make_shared<DefaultAllocator>();
auto allocator = SafeAllocator::FactoryDefaultAllocator();
auto max_degree = this->graph_->MaximumDegree();
auto old_count = this->graph_->TotalCount();
UnorderedMap<InnerIdType, std::shared_ptr<Vector<InnerIdType>>> maps(allocator.get());
Expand Down
2 changes: 1 addition & 1 deletion src/data_cell/sparse_graph_datacell.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ namespace vsag {

SparseGraphDataCell::SparseGraphDataCell(const JsonType& graph_param,
const IndexCommonParam& common_param)
: allocator_(common_param.allocator_), neighbors_(common_param.allocator_) {
: allocator_(common_param.allocator_.get()), neighbors_(common_param.allocator_.get()) {
if (graph_param.contains(GRAPH_PARAM_MAX_DEGREE)) {
this->maximum_degree_ = graph_param[GRAPH_PARAM_MAX_DEGREE];
}
Expand Down
5 changes: 3 additions & 2 deletions src/data_cell/sparse_graph_datacell_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include "default_allocator.h"
#include "fmt/format-inl.h"
#include "graph_interface_test.h"
#include "safe_allocator.h"

using namespace vsag;

Expand All @@ -35,7 +36,7 @@ TestSparseGraphDataCell(const JsonType& graph_param, const IndexCommonParam& par
}

TEST_CASE("graph basic test", "[ut][sparse_graph_datacell]") {
auto allocator = std::make_shared<DefaultAllocator>();
auto allocator = SafeAllocator::FactoryDefaultAllocator();
auto dims = {32, 64};
auto max_degrees = {5, 12, 24, 32, 64, 128};
auto max_capacities = {1, 100, 10000, 10'000'000, 32'179'837};
Expand All @@ -55,7 +56,7 @@ TEST_CASE("graph basic test", "[ut][sparse_graph_datacell]") {
for (auto dim : dims) {
IndexCommonParam param;
param.dim_ = dim;
param.allocator_ = allocator.get();
param.allocator_ = allocator;
for (auto& gp : graph_params) {
TestSparseGraphDataCell(gp, param);
}
Expand Down
10 changes: 0 additions & 10 deletions src/default_allocator.h
Original file line number Diff line number Diff line change
Expand Up @@ -25,16 +25,6 @@
namespace vsag {

class DefaultAllocator : public Allocator {
public:
static std::shared_ptr<Allocator>
Instance() {
static std::shared_ptr<Allocator> s_instance;
if (s_instance == nullptr) {
s_instance = std::make_shared<DefaultAllocator>();
}
return s_instance;
}

public:
DefaultAllocator() = default;
~DefaultAllocator() override {
Expand Down
Loading
Loading