device_options.h 1.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
/**
 * \file lite/load_and_run/src/options/device_options.h
 *
 * This file is part of MegEngine, a deep learning framework developed by
 * Megvii.
 *
 * \copyright Copyright (c) 2020-2021 Megvii Inc. All rights reserved.
 */

#pragma once
#include <gflags/gflags.h>
#include "models/model.h"
#include "option_base.h"

DECLARE_bool(cpu);
#if MGE_WITH_CUDA
DECLARE_bool(cuda);
#endif
DECLARE_bool(cpu_default);
DECLARE_int32(multithread);
DECLARE_int32(multithread_default);
DECLARE_string(multi_thread_core_ids);
namespace lar {

class XPUDeviceOption final : public OptionBase {
public:
    static bool is_valid();
    static std::shared_ptr<OptionBase> create_option();
    void config_model(
            RuntimeParam& runtime_param, std::shared_ptr<ModelBase> model) override;
    std::string option_name() const override { return m_option_name; };

private:
    XPUDeviceOption();
    template <typename ModelImpl>
    void config_model_internel(RuntimeParam&, std::shared_ptr<ModelImpl>){};
    bool enable_cpu;
#if MGE_WITH_CUDA
    bool enable_cuda;
#endif
    bool enable_cpu_default;
    bool enable_multithread;
    bool enable_multithread_default;
    bool enable_set_core_ids;
    size_t thread_num;
    std::vector<int> core_ids;
    std::string m_option_name;
};
}  // namespace lar