1 /**
2 * Copyright 2022 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "src/common/storage.h"
18 #include <sys/stat.h>
19 #ifndef _MSC_VER
20 #include <unistd.h>
21 #endif
22 #include "flatbuffers/flatbuffers.h"
23 #include "src/common/log_adapter.h"
24 #include "src/common/file_utils.h"
25
26 namespace mindspore {
27 namespace lite {
28 namespace {
29 constexpr size_t kMaxNum1024 = 1024;
30 }
Save(const schema::MetaGraphT & graph,const std::string & outputPath)31 int Storage::Save(const schema::MetaGraphT &graph, const std::string &outputPath) {
32 flatbuffers::FlatBufferBuilder builder(kMaxNum1024);
33 auto offset = schema::MetaGraph::Pack(builder, &graph);
34 builder.Finish(offset);
35 schema::FinishMetaGraphBuffer(builder, offset);
36 int size = builder.GetSize();
37 auto content = builder.GetBufferPointer();
38 if (content == nullptr) {
39 MS_LOG(ERROR) << "GetBufferPointer nullptr";
40 return RET_ERROR;
41 }
42 std::string filename = outputPath;
43 if (filename.length() == 0) {
44 MS_LOG(ERROR) << "Invalid output path.";
45 return RET_ERROR;
46 }
47 if (filename.substr(filename.find_last_of(".") + 1) != "ms") {
48 filename = filename + ".ms";
49 }
50 #ifndef _MSC_VER
51 if (access(filename.c_str(), F_OK) == 0) {
52 chmod(filename.c_str(), S_IWUSR);
53 }
54 #endif
55 std::ofstream output(filename, std::ofstream::binary);
56 if (!output.is_open()) {
57 MS_LOG(ERROR) << "Can not open output file: " << filename;
58 return RET_ERROR;
59 }
60 output.write((const char *)content, size);
61 if (output.bad()) {
62 output.close();
63 MS_LOG(ERROR) << "Write output file : " << filename << " failed";
64 return RET_ERROR;
65 }
66 output.close();
67 #ifndef _MSC_VER
68 chmod(filename.c_str(), S_IRUSR);
69 #endif
70 return RET_OK;
71 }
72
Load(const std::string & inputPath)73 schema::MetaGraphT *Storage::Load(const std::string &inputPath) {
74 size_t size = 0;
75 std::string filename = inputPath;
76 if (filename.length() == 0) {
77 MS_LOG(ERROR) << "Invalid input path.";
78 return nullptr;
79 }
80 if (filename.substr(filename.find_last_of(".") + 1) != "ms") {
81 filename = filename + ".ms";
82 }
83 auto buf = ReadFile(filename.c_str(), &size);
84 if (buf == nullptr) {
85 MS_LOG(ERROR) << "the file buffer is nullptr";
86 return nullptr;
87 }
88
89 flatbuffers::Verifier verify((const uint8_t *)buf, size);
90 if (!schema::VerifyMetaGraphBuffer(verify)) {
91 MS_LOG(ERROR) << "the buffer is invalid and fail to create meta graph";
92 return nullptr;
93 }
94
95 auto graphDefT = schema::UnPackMetaGraph(buf);
96 return graphDefT.release();
97 }
98 } // namespace lite
99 } // namespace mindspore
100