option_manager.h 2.7 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
/**
 * \file imperative/src/impl/interpreter/option_manager.h
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#pragma once

#include <string>
#include <unordered_map>

#include "megbrain/common.h"

namespace mgb::imperative::interpreter::intl {

struct OptionManager {
private:
23
    std::unordered_map<std::string, size_t*> m_option_map = {};
24 25
public:
#define DEF_OPTION(name, env_key, default_value, desc) \
26
    size_t name = (m_option_map[#name]=&name, get_option_from_env(env_key, default_value));
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41

    DEF_OPTION(async_level,             "MEGENGINE_INTERP_ASYNC_LEVEL",     2,
        "config whether raise error exactly when invoking op.\n"
        "level 2: both device and user side errors are async;\n"
        "level 1: user side errors are sync;\n"
        "level 0: both sync.");
    DEF_OPTION(enable_swap,             "MEGENGINE_ENABLE_SWAP",            0, "");
    DEF_OPTION(enable_drop,             "MEGENGINE_ENABLE_DROP",            0, "");
    DEF_OPTION(max_recompute_time,      "MEGENGINE_MAX_RECOMP_TIME",        1, "");
    DEF_OPTION(catch_worker_execption,  "MEGENGINE_CATCH_WORKER_EXEC",      1,
        "catch worker exception if enabled, close it when debugging");
    DEF_OPTION(buffer_length,           "MEGENGINE_COMMAND_BUFFER_LENGTH",  3,
        "set command buffer length.");
    DEF_OPTION(enable_host_compute,     "MEGENGINE_HOST_COMPUTE",           1,
        "enable host compute, thus computation may be done in host event if it's device is gpu.");
42 43 44 45 46 47
    DEF_OPTION(enable_auto_drop,        "MEGENGINE_AUTO_DROP",              0, "");
    DEF_OPTION(memory_budget,           "MEGENGINE_MEMORY_BUDGET",          0,
        "auto drop will start whenever gpu memory usage exceeds this value.");
    DEF_OPTION(tensor_lowerbound,       "MEGENGINE_TENSOR_LOWERBOUND",      1048576,
        "the minimum memory value of a tensor added to the candidate set");
    DEF_OPTION(record_computing_path,   "MEGENGINE_RECORD_COMPUTING_PATH",  0, "");
48 49 50

#undef DEF_OPTION

51
    void set_option(const std::string& name, size_t value) {
52 53 54
        *m_option_map[name] = value;
    }

55
    size_t get_option(const std::string& name) const {
56 57 58
        return *m_option_map.at(name);
    }

59
    static size_t get_option_from_env(const std::string& name, size_t default_value) {
60
        if (const char* env_val = MGB_GETENV(name.c_str())) {
61
            sscanf(env_val, "%zu", &default_value);
62 63 64 65 66 67
        }
        return default_value;
    }
};

}