From 10d50d2fb768ae731de6511018aac60a9fd4083d Mon Sep 17 00:00:00 2001 From: fishpenguin Date: Sat, 30 Nov 2019 14:57:35 +0800 Subject: [PATCH] IVF_PQ search on CPUs when using GPU-version --- core/src/scheduler/SchedInst.h | 4 + .../scheduler/optimizer/FaissIVFPQPass.cpp | 74 +++++++++++++++++++ core/src/scheduler/optimizer/FaissIVFPQPass.h | 58 +++++++++++++++ 3 files changed, 136 insertions(+) create mode 100644 core/src/scheduler/optimizer/FaissIVFPQPass.cpp create mode 100644 core/src/scheduler/optimizer/FaissIVFPQPass.h diff --git a/core/src/scheduler/SchedInst.h b/core/src/scheduler/SchedInst.h index 1e8a7acf2e..6cca377033 100644 --- a/core/src/scheduler/SchedInst.h +++ b/core/src/scheduler/SchedInst.h @@ -25,6 +25,7 @@ #include "optimizer/BuildIndexPass.h" #include "optimizer/FaissFlatPass.h" #include "optimizer/FaissIVFFlatPass.h" +#include "optimizer/FaissIVFPQPass.h" #include "optimizer/FaissIVFSQ8HPass.h" #include "optimizer/FaissIVFSQ8Pass.h" #include "optimizer/FallbackPass.h" @@ -129,7 +130,10 @@ class OptimizerInst { pass_list.push_back(std::make_shared()); pass_list.push_back(std::make_shared()); pass_list.push_back(std::make_shared()); +#ifdef CUSTOMIZATION pass_list.push_back(std::make_shared()); +#endif + pass_list.push_back(std::make_shared()); } #endif pass_list.push_back(std::make_shared()); diff --git a/core/src/scheduler/optimizer/FaissIVFPQPass.cpp b/core/src/scheduler/optimizer/FaissIVFPQPass.cpp new file mode 100644 index 0000000000..f97fec63b4 --- /dev/null +++ b/core/src/scheduler/optimizer/FaissIVFPQPass.cpp @@ -0,0 +1,74 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. +#ifdef MILVUS_GPU_VERSION +#include "scheduler/optimizer/FaissIVFPQPass.h" +#include "cache/GpuCacheMgr.h" +#include "scheduler/SchedInst.h" +#include "scheduler/Utils.h" +#include "scheduler/task/SearchTask.h" +#include "scheduler/tasklabel/SpecResLabel.h" +#include "server/Config.h" +#include "utils/Log.h" + +namespace milvus { +namespace scheduler { + +void +FaissIVFPQPass::Init() { +#ifdef MILVUS_GPU_VERSION + server::Config& config = server::Config::GetInstance(); + Status s = config.GetEngineConfigGpuSearchThreshold(threshold_); + if (!s.ok()) { + threshold_ = std::numeric_limits::max(); + } + s = config.GetGpuResourceConfigSearchResources(gpus); + if (!s.ok()) { + throw; + } +#endif +} + +bool +FaissIVFPQPass::Run(const TaskPtr& task) { + if (task->Type() != TaskType::SearchTask) { + return false; + } + + auto search_task = std::static_pointer_cast(task); + if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_PQ) { + return false; + } + + auto search_job = std::static_pointer_cast(search_task->job_.lock()); + ResourcePtr res_ptr; + if (search_job->nq() < threshold_) { + SERVER_LOG_DEBUG << "FaissIVFPQPass: nq < gpu_search_threshold, specify cpu to search!"; + res_ptr = ResMgrInst::GetInstance()->GetResource("cpu"); + } else { + auto best_device_id = count_ % gpus.size(); + SERVER_LOG_DEBUG << "FaissIVFPQPass: nq > gpu_search_threshold, specify gpu" << best_device_id << " to search!"; + count_++; + res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, gpus[best_device_id]); + } + auto label = std::make_shared(res_ptr); + task->label() = label; + return true; +} + +} // namespace scheduler +} // namespace milvus +#endif diff --git a/core/src/scheduler/optimizer/FaissIVFPQPass.h b/core/src/scheduler/optimizer/FaissIVFPQPass.h new file mode 100644 index 0000000000..9225f84b7c --- /dev/null +++ b/core/src/scheduler/optimizer/FaissIVFPQPass.h @@ -0,0 +1,58 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. +#ifdef MILVUS_GPU_VERSION +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "Pass.h" + +namespace milvus { +namespace scheduler { + +class FaissIVFPQPass : public Pass { + public: + FaissIVFPQPass() = default; + + public: + void + Init() override; + + bool + Run(const TaskPtr& task) override; + + private: + int64_t threshold_ = std::numeric_limits::max(); + int64_t count_ = 0; + std::vector gpus; +}; + +using FaissIVFPQPassPtr = std::shared_ptr; + +} // namespace scheduler +} // namespace milvus +#endif