forked from ROCm/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
module.h
71 lines (60 loc) · 2.42 KB
/
module.h
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
/**
* A global dictionary that holds information about what Caffe2 modules have
* been loaded in the current runtime, and also utility functions to load
* modules.
*/
#ifndef CAFFE2_CORE_MODULE_H_
#define CAFFE2_CORE_MODULE_H_
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <memory>
#include <mutex>
#include "caffe2/core/common.h"
#include <c10/util/typeid.h>
namespace caffe2 {
/**
* A module schema that can be used to store specific information about
* different modules. Currently, we only store the name and a simple
* description of what this module does.
*/
class TORCH_API ModuleSchema {
public:
ModuleSchema(const char* name, const char* description);
};
/**
* @brief Current Modules present in the Caffe2 runtime.
* Returns:
* map: a map of modules and (optionally) their description. The key is the
* module name, and the value is the description for that module. The
* module name is recommended to be the part that constitutes the trunk
* of the dynamic library: for example, a module called
* libcaffe2_db_rocksdb.so should have the name "caffe2_db_rocksdb". The
* reason we do not use "lib" is because it's somewhat redundant, and
* the reason we do not include ".so" is for cross-platform compatibility
* on platforms like mac os.
*/
TORCH_API const CaffeMap<string, const ModuleSchema*>& CurrentModules();
/**
* @brief Checks whether a module is already present in the current binary.
*/
TORCH_API bool HasModule(const string& name);
/**
* @brief Load a module.
* Inputs:
* name: a module name or a path name.
* It is recommended that you use the name of the module, and leave the
* full path option to only experimental modules.
* filename: (optional) a filename that serves as a hint to load the module.
*/
TORCH_API void LoadModule(const string& name, const string& filename="");
#define CAFFE2_MODULE(name, description) \
extern "C" { \
bool gCaffe2ModuleSanityCheck##name() { return true; } \
} \
namespace { \
static ::caffe2::ModuleSchema module_schema_##name(#name, description); \
}
} // namespace caffe2
#endif // CAFFE2_CORE_MODULE_H_