Skip to content

Commit bb66af0

Browse files
authored
Refactor data loader creation. (#13208)
1 parent 50ffe44 commit bb66af0

File tree

1 file changed

+11
-23
lines changed

1 file changed

+11
-23
lines changed

extension/module/module.cpp

Lines changed: 11 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -42,28 +42,28 @@ using ET_RUNTIME_NAMESPACE::MethodMeta;
4242
using ET_RUNTIME_NAMESPACE::Program;
4343

4444
namespace {
45-
runtime::Result<std::unique_ptr<runtime::DataLoader>> load_file(
45+
runtime::Result<std::unique_ptr<runtime::DataLoader>> make_data_loader(
4646
const std::string& file_path,
4747
Module::LoadMode mode) {
48-
std::unique_ptr<runtime::DataLoader> res = nullptr;
48+
std::unique_ptr<runtime::DataLoader> data_loader;
4949
switch (mode) {
5050
case Module::LoadMode::File:
51-
res = ET_UNWRAP_UNIQUE(FileDataLoader::from(file_path.c_str()));
51+
data_loader = ET_UNWRAP_UNIQUE(FileDataLoader::from(file_path.c_str()));
5252
break;
5353
case Module::LoadMode::Mmap:
54-
res = ET_UNWRAP_UNIQUE(MmapDataLoader::from(
54+
data_loader = ET_UNWRAP_UNIQUE(MmapDataLoader::from(
5555
file_path.c_str(), MmapDataLoader::MlockConfig::NoMlock));
5656
break;
5757
case Module::LoadMode::MmapUseMlock:
58-
res = ET_UNWRAP_UNIQUE(MmapDataLoader::from(file_path.c_str()));
58+
data_loader = ET_UNWRAP_UNIQUE(MmapDataLoader::from(file_path.c_str()));
5959
break;
6060
case Module::LoadMode::MmapUseMlockIgnoreErrors:
61-
res = ET_UNWRAP_UNIQUE(MmapDataLoader::from(
61+
data_loader = ET_UNWRAP_UNIQUE(MmapDataLoader::from(
6262
file_path.c_str(),
6363
MmapDataLoader::MlockConfig::UseMlockIgnoreErrors));
6464
break;
6565
}
66-
return res;
66+
return data_loader;
6767
}
6868
} // namespace
6969

@@ -137,29 +137,17 @@ Module::Module(
137137

138138
runtime::Error Module::load(const Program::Verification verification) {
139139
if (!is_loaded()) {
140-
// Load the program
141140
if (!data_loader_) {
142-
auto res = load_file(file_path_, load_mode_);
143-
if (!res.ok()) {
144-
return res.error();
145-
}
146-
data_loader_ = std::move(res.get());
141+
data_loader_ = ET_UNWRAP(make_data_loader(file_path_, load_mode_));
147142
}
148-
// If a .ptd path was given load it.
149-
if (data_map_path_ != "") {
150-
auto res = load_file(data_map_path_, load_mode_);
151-
if (!res.ok()) {
152-
return res.error();
153-
}
154-
data_map_loader_ = std::move(res.get());
143+
if (!data_map_path_.empty()) {
144+
data_map_loader_ =
145+
ET_UNWRAP(make_data_loader(data_map_path_, load_mode_));
155146
}
156-
// If we have a .ptd loader, then load the map.
157147
if (data_map_loader_) {
158148
data_map_ =
159149
ET_UNWRAP_UNIQUE(FlatTensorDataMap::load(data_map_loader_.get()));
160150
}
161-
// else: either the map itself was provided or we have no data map, either
162-
// way no work to do.
163151
auto program =
164152
ET_UNWRAP_UNIQUE(Program::load(data_loader_.get(), verification));
165153
program_ = std::shared_ptr<Program>(

0 commit comments

Comments
 (0)