mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
fb18c29486
s/Facebook, Inc./Meta Platforms, Inc/ s/Confidential and proprietary./This source code is licensed under the BSD-style license/ Per https://www.internalfb.com/intern/wiki/Open_Source/Licenses/Straight_BSD/ Also, add linter that prevents adding those in the future Fixes https://github.com/pytorch/pytorch/issues/90187 Pull Request resolved: https://github.com/pytorch/pytorch/pull/90805 Approved by: https://github.com/zpao
77 lines
2.3 KiB
C++
77 lines
2.3 KiB
C++
// Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
//
|
|
// This is a very basic knob implementation that purely uses command line flags.
|
|
// This can be replaced with a more sophisticated implementation for use in
|
|
// other production environments.
|
|
|
|
#include <map>
|
|
|
|
#include <c10/util/string_view.h>
|
|
#include <c10/util/Flags.h>
|
|
|
|
#include "caffe2/utils/knobs.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
namespace detail {
|
|
// Get the map of knob names to pointers to their command-line controlled
|
|
// boolean value.
|
|
std::map<c10::string_view, bool*>& getRegisteredKnobs() {
|
|
// It's safe to store the keys as string_view, since DEFINE_KNOB() ensures
|
|
// that these views always point to string literals.
|
|
static std::map<c10::string_view, bool*> registeredKnobs;
|
|
return registeredKnobs;
|
|
}
|
|
} // namespace detail
|
|
|
|
bool CheckKnob(c10::string_view name) {
|
|
const auto& knobs = detail::getRegisteredKnobs();
|
|
auto iter = knobs.find(name);
|
|
if (iter == knobs.end()) {
|
|
throw std::invalid_argument(
|
|
"attempted to check unknown knob \"" + std::string(name) + "\"");
|
|
}
|
|
return *iter->second;
|
|
}
|
|
|
|
namespace {
|
|
class RegisterKnob {
|
|
public:
|
|
RegisterKnob(c10::string_view name, bool* cmdlineFlag) {
|
|
auto ret = caffe2::detail::getRegisteredKnobs().emplace(name, cmdlineFlag);
|
|
if (!ret.second) {
|
|
throw std::runtime_error("duplicate knob name: " + std::string(name));
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
} // namespace caffe2
|
|
|
|
/**
|
|
* Define a knob.
|
|
*
|
|
* This will define a --caffe2_knob_<name> command line flag to control the
|
|
* knob.
|
|
*
|
|
* The knob can be checked in code by calling CheckKnob(name)
|
|
* or CheckKnob<check_fn_name>()
|
|
*/
|
|
#define DEFINE_KNOB(name, check_fn_name, default_value, docstring) \
|
|
C10_DEFINE_bool(caffe2_knob_##name, default_value, docstring); \
|
|
namespace caffe2 { \
|
|
bool CheckKnob##check_fn_name() { \
|
|
return FLAGS_caffe2_knob_##name; \
|
|
} \
|
|
} \
|
|
static caffe2::RegisterKnob _knob_##name(#name, &FLAGS_caffe2_knob_##name)
|
|
|
|
/*
|
|
* Definitions of well-known knobs.
|
|
*/
|
|
|
|
DEFINE_KNOB(
|
|
example_knob,
|
|
ExampleKnob,
|
|
false,
|
|
"An example knob, mainly intended for use in unit tests");
|