10
10
#include < cassert>
11
11
#include < cstdlib>
12
12
#include < sstream>
13
+ #include < regex>
14
+ #include < iterator>
13
15
#ifdef _MSC_VER
14
16
#include < intrin.h>
15
17
#endif
@@ -75,12 +77,20 @@ bool LLModel::Implementation::isImplementation(const Dlhandle &dl) {
75
77
return dl.get <bool (uint32_t )>(" is_g4a_backend_model_implementation" );
76
78
}
77
79
78
- const std::vector<LLModel::Implementation> &LLModel::Implementation::implementationList () {
80
+ static std::vector<std::unique_ptr<LLModel::Implementation>> s_impl_libs;
81
+ static bool s_scanned = false ;
82
+
83
+ const std::vector<std::unique_ptr<LLModel::Implementation>> &LLModel::Implementation::implementationList () {
79
84
// NOTE: allocated on heap so we leak intentionally on exit so we have a chance to clean up the
80
85
// individual models without the cleanup of the static list interfering
81
- static auto * libs = new std::vector<Implementation>([] () {
82
- std::vector<Implementation> fres;
83
-
86
+ if (!s_scanned) {
87
+ std::string impl_name_re = " (bert|llama|gptj|llamamodel-mainline)" ;
88
+ if (requires_avxonly ()) {
89
+ impl_name_re += " -avxonly" ;
90
+ } else {
91
+ impl_name_re += " -default" ;
92
+ }
93
+ std::regex re (impl_name_re);
84
94
auto search_in_directory = [&](const std::string& paths) {
85
95
std::stringstream ss (paths);
86
96
std::string path;
@@ -90,32 +100,34 @@ const std::vector<LLModel::Implementation> &LLModel::Implementation::implementat
90
100
// Iterate over all libraries
91
101
for (const auto & f : std::filesystem::directory_iterator (fs_path)) {
92
102
const std::filesystem::path& p = f.path ();
103
+
93
104
if (p.extension () != LIB_FILE_EXT) continue ;
105
+ if (!std::regex_search (p.stem ().string (), re)) continue ;
106
+
94
107
// Add to list if model implementation
95
108
try {
109
+ std::cerr << " attempt load: " << p << std::endl;
96
110
Dlhandle dl (p.string ());
97
111
if (!Implementation::isImplementation (dl)) {
98
112
continue ;
99
113
}
100
- fres .emplace_back (Implementation ( std::move (dl)));
114
+ s_impl_libs .emplace_back (std::make_unique< Implementation>( Implementation ( std::move (dl) )));
101
115
} catch (...) {}
102
116
}
103
117
}
104
118
};
105
119
106
120
search_in_directory (s_implementations_search_path);
107
-
108
- return fres;
109
- }());
110
- // Return static result
111
- return *libs;
121
+ s_scanned = true ;
122
+ };
123
+ return s_impl_libs;
112
124
}
113
125
114
126
const LLModel::Implementation* LLModel::Implementation::implementation (const char *fname, const std::string& buildVariant) {
115
127
for (const auto & i : implementationList ()) {
116
- if (buildVariant != i. m_buildVariant ) continue ;
117
- if (!i. m_magicMatch (fname)) continue ;
118
- return &i ;
128
+ if (buildVariant != i-> m_buildVariant ) continue ;
129
+ if (!i-> m_magicMatch (fname)) continue ;
130
+ return i. get () ;
119
131
}
120
132
return nullptr ;
121
133
}
@@ -170,6 +182,7 @@ LLModel *LLModel::Implementation::construct(const std::string &modelPath, std::s
170
182
171
183
void LLModel::Implementation::setImplementationsSearchPath (const std::string& path) {
172
184
s_implementations_search_path = path;
185
+ s_scanned = false ;
173
186
}
174
187
175
188
const std::string& LLModel::Implementation::implementationsSearchPath () {
0 commit comments