1 /**
2 * Copyright 2020 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 "minddata/dataset/engine/gnn/graph_shared_memory.h"
18
19 #include <string>
20 #include "utils/file_utils.h"
21 #include "utils/ms_utils.h"
22 #include "minddata/dataset/util/log_adapter.h"
23
24 namespace mindspore {
25 namespace dataset {
26 namespace gnn {
27
GraphSharedMemory(int64_t memory_size,key_t memory_key)28 GraphSharedMemory::GraphSharedMemory(int64_t memory_size, key_t memory_key)
29 : memory_size_(memory_size),
30 memory_key_(memory_key),
31 memory_ptr_(nullptr),
32 memory_offset_(0),
33 is_new_create_(false) {
34 std::stringstream stream;
35 stream << std::hex << memory_key_;
36 memory_key_str_ = stream.str();
37 }
38
GraphSharedMemory(int64_t memory_size,const std::string & mr_file)39 GraphSharedMemory::GraphSharedMemory(int64_t memory_size, const std::string &mr_file)
40 : mr_file_(mr_file),
41 memory_size_(memory_size),
42 memory_key_(-1),
43 memory_ptr_(nullptr),
44 memory_offset_(0),
45 is_new_create_(false) {}
46
~GraphSharedMemory()47 GraphSharedMemory::~GraphSharedMemory() {
48 if (is_new_create_) {
49 (void)DeleteSharedMemory();
50 }
51 }
52
CreateSharedMemory()53 Status GraphSharedMemory::CreateSharedMemory() {
54 if (memory_key_ == -1) {
55 // ftok to generate unique key
56 auto realpath = FileUtils::GetRealPath(mr_file_.data());
57 CHECK_FAIL_RETURN_UNEXPECTED(realpath.has_value(), "Get real path failed, path=" + mr_file_);
58 memory_key_ = ftok(common::SafeCStr(realpath.value()), kGnnSharedMemoryId);
59 CHECK_FAIL_RETURN_UNEXPECTED(memory_key_ != -1, "Failed to get key of shared memory. file_name:" + mr_file_);
60 std::stringstream stream;
61 stream << std::hex << memory_key_;
62 memory_key_str_ = stream.str();
63 }
64 int shmflg = (0666 | IPC_CREAT | IPC_EXCL);
65 Status s = SharedMemoryImpl(shmflg);
66 if (s.IsOk()) {
67 is_new_create_ = true;
68 MS_LOG(INFO) << "Create shared memory success, key=0x" << memory_key_str_;
69 } else {
70 MS_LOG(WARNING) << "Shared memory with the same key may already exist, key=0x" << memory_key_str_;
71 shmflg = (0666 | IPC_CREAT);
72 s = SharedMemoryImpl(shmflg);
73 if (!s.IsOk()) {
74 RETURN_STATUS_UNEXPECTED("Create shared memory fao;ed, key=0x" + memory_key_str_);
75 }
76 }
77 return Status::OK();
78 }
79
GetSharedMemory()80 Status GraphSharedMemory::GetSharedMemory() {
81 int shmflg = 0;
82 RETURN_IF_NOT_OK(SharedMemoryImpl(shmflg));
83 return Status::OK();
84 }
85
DeleteSharedMemory()86 Status GraphSharedMemory::DeleteSharedMemory() {
87 int shmid = shmget(memory_key_, 0, 0);
88 CHECK_FAIL_RETURN_UNEXPECTED(shmid != -1, "Failed to get shared memory. key=0x" + memory_key_str_);
89 int result = shmctl(shmid, IPC_RMID, 0);
90 CHECK_FAIL_RETURN_UNEXPECTED(result != -1, "Failed to delete shared memory. key=0x" + memory_key_str_);
91 return Status::OK();
92 }
93
SharedMemoryImpl(const int & shmflg)94 Status GraphSharedMemory::SharedMemoryImpl(const int &shmflg) {
95 // shmget returns an identifier in shmid
96 CHECK_FAIL_RETURN_UNEXPECTED(memory_size_ >= 0, "Invalid memory size, should be greater than zero.");
97 int shmid = shmget(memory_key_, memory_size_, shmflg);
98 CHECK_FAIL_RETURN_UNEXPECTED(shmid != -1, "Failed to get shared memory. key=0x" + memory_key_str_);
99
100 // shmat to attach to shared memory
101 auto data = shmat(shmid, reinterpret_cast<void *>(0), 0);
102 CHECK_FAIL_RETURN_UNEXPECTED(data != (char *)(-1), "Failed to address shared memory. key=0x" + memory_key_str_);
103 memory_ptr_ = reinterpret_cast<uint8_t *>(data);
104
105 return Status::OK();
106 }
107
InsertData(const uint8_t * data,int64_t len,int64_t * offset)108 Status GraphSharedMemory::InsertData(const uint8_t *data, int64_t len, int64_t *offset) {
109 CHECK_FAIL_RETURN_UNEXPECTED(data, "Input data is nullptr.");
110 CHECK_FAIL_RETURN_UNEXPECTED(len > 0, "Input len is invalid.");
111 CHECK_FAIL_RETURN_UNEXPECTED(offset, "Input offset is nullptr.");
112
113 std::lock_guard<std::mutex> lck(mutex_);
114 CHECK_FAIL_RETURN_UNEXPECTED((memory_size_ - memory_offset_ >= len),
115 "Insufficient shared memory space to insert data.");
116 if (EOK != memcpy_s(memory_ptr_ + memory_offset_, memory_size_ - memory_offset_, data, len)) {
117 RETURN_STATUS_UNEXPECTED("Failed to insert data into shared memory.");
118 }
119 *offset = memory_offset_;
120 memory_offset_ += len;
121 return Status::OK();
122 }
123
GetData(uint8_t * data,int64_t data_len,int64_t offset,int64_t get_data_len)124 Status GraphSharedMemory::GetData(uint8_t *data, int64_t data_len, int64_t offset, int64_t get_data_len) {
125 CHECK_FAIL_RETURN_UNEXPECTED(data, "Input data is nullptr.");
126 CHECK_FAIL_RETURN_UNEXPECTED(get_data_len > 0, "Input get_data_len is invalid.");
127 CHECK_FAIL_RETURN_UNEXPECTED(data_len >= get_data_len, "Insufficient target address space.");
128
129 CHECK_FAIL_RETURN_UNEXPECTED(memory_size_ >= get_data_len + offset,
130 "get_data_len is too large, beyond the space of shared memory.");
131 if (EOK != memcpy_s(data, data_len, memory_ptr_ + offset, get_data_len)) {
132 RETURN_STATUS_UNEXPECTED("Failed to insert data into shared memory.");
133 }
134 return Status::OK();
135 }
136
137 } // namespace gnn
138 } // namespace dataset
139 } // namespace mindspore
140