MCPcopy Create free account
hub / github.com/MegEngine/MegCC / alignment_or_alloc_weights

Function alignment_or_alloc_weights

runtime/src/parse.c:223–295  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

221}
222
223static TinyNNStatus alignment_or_alloc_weights(
224 CombineModel* model, void* buffer, int share_weights) {
225 int nr_weight = model->nr_origin_weight;
226 uintptr_t current_addr = (uintptr_t)buffer;
227 Device* host_dev = &model->host_dev;
228 unsigned int alignment = host_dev->alignment;
229 //! sort all the weights by their ptr
230 Tensor** sorted_weights = tinynn_malloc(sizeof(Tensor*) * nr_weight);
231 for (int i = 0; i < nr_weight; i++) {
232 sorted_weights[i] = model->weights + i;
233 }
234 qsort(sorted_weights, nr_weight, sizeof(Tensor*), weight_ptr_compare);
235 //! align or alloc them
236 for (int i = 0; i < nr_weight; i++) {
237 Tensor* weight = sorted_weights[i];
238 size_t length = weight->size;
239 size_t layout_size = tensor_length_in_byte(weight);
240 int compressed = 0;
241 if (length > 0 && weight->dtype.type_enum == TinyNN_FLOAT &&
242 length == layout_size / 2) {
243 compressed = 1;
244 LOG_DEBUG(
245 "The weight data is compressed, decompressed it to "
246 "fp32.\n");
247 }
248
249 void* origin_ptr = weight->ptr;
250 TINYNN_ASSERT_MSG(
251 current_addr < (uintptr_t)origin_ptr, "weight pointer addr error.\n");
252 uintptr_t alignment_addr =
253 (uintptr_t)origin_ptr & (~((uintptr_t)alignment - 1));
254 if (!share_weights || alignment_addr < current_addr || compressed) {
255 if (!(weight->ptr = host_dev->malloc(layout_size))) {
256 LOG_ERROR("malloc weight memory fail.\n");
257 }
258 weight->is_shared = 0;
259 if (!compressed) {
260 memcpy(weight->ptr, origin_ptr, length);
261 } else {
262 weight->size = layout_size;
263 uint16_t* src = (uint16_t*)origin_ptr;
264 float* dst = (float*)(weight->ptr);
265 for (int id = 0; id < layout_size / sizeof(float); id++) {
266 *dst = half_to_float(*src);
267 dst++;
268 src++;
269 }
270 }
271 } else {
272 uintptr_t more_fit_addr =
273 (current_addr + alignment - 1) & (~((uintptr_t)alignment - 1));
274 alignment_addr =
275 alignment_addr > more_fit_addr ? more_fit_addr : alignment_addr;
276 if (alignment_addr != (uintptr_t)origin_ptr) {
277 LOG_DEBUG(
278 "align weight: %d from %p to %p\n", i, origin_ptr,
279 (void*)alignment_addr);
280 TINYNN_ASSERT_MSG(

Callers 1

parse_modelFunction · 0.85

Calls 5

tinynn_mallocFunction · 0.85
tensor_length_in_byteFunction · 0.85
half_to_floatFunction · 0.85
tinynn_freeFunction · 0.85
mallocMethod · 0.80

Tested by

no test coverage detected