1 // Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //     http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 /*-------------------------------------------
16                 Includes
17 -------------------------------------------*/
18 #include "rk_mpi_mmz.h"
19 #include "rknn_api.h"
20 
21 #include <float.h>
22 #include <stdio.h>
23 #include <stdlib.h>
24 #include <string.h>
25 #include <sys/time.h>
26 
27 #define STB_IMAGE_IMPLEMENTATION
28 #include "stb/stb_image.h"
29 #define STB_IMAGE_RESIZE_IMPLEMENTATION
30 #include <stb/stb_image_resize.h>
31 
32 /*-------------------------------------------
33                   Functions
34 -------------------------------------------*/
getCurrentTimeUs()35 static inline int64_t getCurrentTimeUs()
36 {
37   struct timeval tv;
38   gettimeofday(&tv, NULL);
39   return tv.tv_sec * 1000000 + tv.tv_usec;
40 }
41 
rknn_GetTopN(float * pfProb,float * pfMaxProb,uint32_t * pMaxClass,uint32_t outputCount,uint32_t topNum)42 static int rknn_GetTopN(float* pfProb, float* pfMaxProb, uint32_t* pMaxClass, uint32_t outputCount, uint32_t topNum)
43 {
44   uint32_t i, j;
45   uint32_t top_count = outputCount > topNum ? topNum : outputCount;
46 
47   for (i = 0; i < topNum; ++i) {
48     pfMaxProb[i] = -FLT_MAX;
49     pMaxClass[i] = -1;
50   }
51 
52   for (j = 0; j < top_count; j++) {
53     for (i = 0; i < outputCount; i++) {
54       if ((i == *(pMaxClass + 0)) || (i == *(pMaxClass + 1)) || (i == *(pMaxClass + 2)) || (i == *(pMaxClass + 3)) ||
55           (i == *(pMaxClass + 4))) {
56         continue;
57       }
58 
59       if (pfProb[i] > *(pfMaxProb + j)) {
60         *(pfMaxProb + j) = pfProb[i];
61         *(pMaxClass + j) = i;
62       }
63     }
64   }
65 
66   return 1;
67 }
68 
dump_tensor_attr(rknn_tensor_attr * attr)69 static void dump_tensor_attr(rknn_tensor_attr* attr)
70 {
71   printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
72          "zp=%d, scale=%f\n",
73          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
74          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
75          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
76 }
77 
load_image(const char * image_path,rknn_tensor_attr * input_attr)78 static unsigned char* load_image(const char* image_path, rknn_tensor_attr* input_attr)
79 {
80   int req_height  = 0;
81   int req_width   = 0;
82   int req_channel = 0;
83 
84   switch (input_attr->fmt) {
85   case RKNN_TENSOR_NHWC:
86     req_height  = input_attr->dims[1];
87     req_width   = input_attr->dims[2];
88     req_channel = input_attr->dims[3];
89     break;
90   case RKNN_TENSOR_NCHW:
91     req_height  = input_attr->dims[2];
92     req_width   = input_attr->dims[3];
93     req_channel = input_attr->dims[1];
94     break;
95   default:
96     printf("meet unsupported layout\n");
97     return NULL;
98   }
99 
100   int height  = 0;
101   int width   = 0;
102   int channel = 0;
103 
104   unsigned char* image_data = stbi_load(image_path, &width, &height, &channel, req_channel);
105   if (image_data == NULL) {
106     printf("load image failed!\n");
107     return NULL;
108   }
109 
110   if (width != req_width || height != req_height) {
111     unsigned char* image_resized = (unsigned char*)STBI_MALLOC(req_width * req_height * req_channel);
112     if (!image_resized) {
113       printf("malloc image failed!\n");
114       STBI_FREE(image_data);
115       return NULL;
116     }
117     if (stbir_resize_uint8(image_data, width, height, 0, image_resized, req_width, req_height, 0, channel) != 1) {
118       printf("resize image failed!\n");
119       STBI_FREE(image_data);
120       return NULL;
121     }
122     STBI_FREE(image_data);
123     image_data = image_resized;
124   }
125 
126   return image_data;
127 }
128 
129 /*-------------------------------------------
130                   Main Functions
131 -------------------------------------------*/
main(int argc,char * argv[])132 int main(int argc, char* argv[])
133 {
134   if (argc < 3) {
135     printf("Usage:%s model_path input_path [loop_count]\n", argv[0]);
136     return -1;
137   }
138 
139   char* model_path = argv[1];
140   char* input_path = argv[2];
141 
142   int loop_count = 1;
143   if (argc > 3) {
144     loop_count = atoi(argv[3]);
145   }
146 
147   rknn_context ctx = 0;
148 
149   // Load RKNN Model
150   int ret = rknn_init(&ctx, model_path, 0, RKNN_FLAG_MEM_ALLOC_OUTSIDE, NULL);
151   if (ret < 0) {
152     printf("rknn_init fail! ret=%d\n", ret);
153     return -1;
154   }
155 
156   // Get sdk and driver version
157   rknn_sdk_version sdk_ver;
158   ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &sdk_ver, sizeof(sdk_ver));
159   if (ret != RKNN_SUCC) {
160     printf("rknn_query fail! ret=%d\n", ret);
161     return -1;
162   }
163   printf("rknn_api/rknnrt version: %s, driver version: %s\n", sdk_ver.api_version, sdk_ver.drv_version);
164 
165   // Get weight and internal mem size
166   rknn_mem_size mem_size;
167   ret = rknn_query(ctx, RKNN_QUERY_MEM_SIZE, &mem_size, sizeof(mem_size));
168   if (ret != RKNN_SUCC) {
169     printf("rknn_query fail! ret=%d\n", ret);
170     return -1;
171   }
172   printf("total weight size: %d, total internal size: %d\n", mem_size.total_weight_size, mem_size.total_internal_size);
173 
174   // Get Model Input Output Info
175   rknn_input_output_num io_num;
176   ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
177   if (ret != RKNN_SUCC) {
178     printf("rknn_query fail! ret=%d\n", ret);
179     return -1;
180   }
181   printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);
182 
183   printf("input tensors:\n");
184   rknn_tensor_attr input_attrs[io_num.n_input];
185   memset(input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
186   for (uint32_t i = 0; i < io_num.n_input; i++) {
187     input_attrs[i].index = i;
188     // query info
189     ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
190     if (ret < 0) {
191       printf("rknn_init error! ret=%d\n", ret);
192       return -1;
193     }
194     dump_tensor_attr(&input_attrs[i]);
195   }
196 
197   printf("output tensors:\n");
198   rknn_tensor_attr output_attrs[io_num.n_output];
199   memset(output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
200   for (uint32_t i = 0; i < io_num.n_output; i++) {
201     output_attrs[i].index = i;
202     // query info
203     ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
204     if (ret != RKNN_SUCC) {
205       printf("rknn_query fail! ret=%d\n", ret);
206       return -1;
207     }
208     dump_tensor_attr(&output_attrs[i]);
209   }
210 
211   // Get custom string
212   rknn_custom_string custom_string;
213   ret = rknn_query(ctx, RKNN_QUERY_CUSTOM_STRING, &custom_string, sizeof(custom_string));
214   if (ret != RKNN_SUCC) {
215     printf("rknn_query fail! ret=%d\n", ret);
216     return -1;
217   }
218   printf("custom string: %s\n", custom_string.string);
219 
220   unsigned char*     input_data   = NULL;
221   rknn_tensor_type   input_type   = RKNN_TENSOR_UINT8;
222   rknn_tensor_format input_layout = RKNN_TENSOR_NHWC;
223 
224   // Load image
225   input_data = load_image(input_path, &input_attrs[0]);
226 
227   if (!input_data) {
228     return -1;
229   }
230 
231   int mb_flags = RK_MMZ_ALLOC_TYPE_CMA | RK_MMZ_ALLOC_UNCACHEABLE;
232 
233   // Allocate weight memory in outside
234   MB_BLK           weight_mb;
235   rknn_tensor_mem* weight_mem;
236   ret = RK_MPI_MMZ_Alloc(&weight_mb, mem_size.total_weight_size, mb_flags);
237   if (ret < 0) {
238     printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
239     return ret;
240   }
241   void* weight_virt = RK_MPI_MMZ_Handle2VirAddr(weight_mb);
242   if (weight_virt == NULL) {
243     printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
244     return -1;
245   }
246   int weight_fd = RK_MPI_MMZ_Handle2Fd(weight_mb);
247   if (weight_fd < 0) {
248     printf("RK_MPI_MMZ_Handle2Fd failed!\n");
249     return -1;
250   }
251   weight_mem = rknn_create_mem_from_fd(ctx, weight_fd, weight_virt, mem_size.total_weight_size, 0);
252   ret        = rknn_set_weight_mem(ctx, weight_mem);
253   if (ret < 0) {
254     printf("rknn_set_weight_mem fail! ret=%d\n", ret);
255     return -1;
256   }
257   printf("weight mb info: virt = %p, fd = %d, size: %d\n", weight_virt, weight_fd, mem_size.total_weight_size);
258 
259   // Allocate internal memory in outside
260   MB_BLK           internal_mb;
261   rknn_tensor_mem* internal_mem;
262   ret = RK_MPI_MMZ_Alloc(&internal_mb, mem_size.total_internal_size, mb_flags);
263   if (ret < 0) {
264     printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
265     return ret;
266   }
267   void* internal_virt = RK_MPI_MMZ_Handle2VirAddr(internal_mb);
268   if (internal_virt == NULL) {
269     printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
270     return -1;
271   }
272   int internal_fd = RK_MPI_MMZ_Handle2Fd(internal_mb);
273   if (internal_fd < 0) {
274     printf("RK_MPI_MMZ_Handle2Fd failed!\n");
275     return -1;
276   }
277   internal_mem = rknn_create_mem_from_fd(ctx, internal_fd, internal_virt, mem_size.total_internal_size, 0);
278   ret          = rknn_set_internal_mem(ctx, internal_mem);
279   if (ret < 0) {
280     printf("rknn_set_internal_mem fail! ret=%d\n", ret);
281     return -1;
282   }
283   printf("internal mb info: virt = %p, fd = %d, size: %d\n", internal_virt, internal_fd, mem_size.total_internal_size);
284 
285   // Allocate input memory in outside
286   MB_BLK input_mb;
287   int    input_size = input_attrs[0].size_with_stride;
288   ret               = RK_MPI_MMZ_Alloc(&input_mb, input_size, mb_flags);
289   if (ret < 0) {
290     printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
291     return ret;
292   }
293   void* input_virt = RK_MPI_MMZ_Handle2VirAddr(input_mb);
294   if (input_virt == NULL) {
295     printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
296     return -1;
297   }
298   int input_fd = RK_MPI_MMZ_Handle2Fd(input_mb);
299   if (input_fd < 0) {
300     printf("RK_MPI_MMZ_Handle2Fd failed!\n");
301     return -1;
302   }
303   printf("input mb info: virt = %p, fd = %d, size: %d\n", input_virt, input_fd, input_size);
304 
305   // Allocate outputs memory in outside
306   MB_BLK output_mbs[io_num.n_output];
307   void*  output_virts[io_num.n_output];
308   int    output_fds[io_num.n_output];
309   for (uint32_t i = 0; i < io_num.n_output; ++i) {
310     // default output type is depend on model, this require float32 to compute top5
311     output_attrs[i].type = RKNN_TENSOR_FLOAT32;
312     int output_size      = output_attrs[i].n_elems * sizeof(float);
313     output_attrs[i].size = output_size;
314     ret                  = RK_MPI_MMZ_Alloc(&output_mbs[i], output_size, mb_flags);
315     if (ret < 0) {
316       printf("RK_MPI_MMZ_Alloc failed, ret: %d\n", ret);
317       return ret;
318     }
319     output_virts[i] = RK_MPI_MMZ_Handle2VirAddr(output_mbs[i]);
320     if (output_virts[i] == NULL) {
321       printf("RK_MPI_MMZ_Handle2VirAddr failed!\n");
322       return -1;
323     }
324     output_fds[i] = RK_MPI_MMZ_Handle2Fd(output_mbs[i]);
325     if (output_fds[i] < 0) {
326       printf("RK_MPI_MMZ_Handle2Fd failed!\n");
327       return -1;
328     }
329     printf("output%d mb info: virt = %p, fd = %d, size = %d\n", i, output_virts[i], output_fds[i], output_size);
330   }
331 
332   // Create input tensor memory
333   rknn_tensor_mem* input_mems[1];
334   // default input type is int8 (normalize and quantize need compute in outside)
335   // if set uint8, will fuse normalize and quantize to npu
336   input_attrs[0].type = input_type;
337   // default fmt is NHWC, npu only support NHWC in zero copy mode
338   input_attrs[0].fmt = input_layout;
339 
340   input_mems[0] = rknn_create_mem_from_fd(ctx, input_fd, input_virt, input_attrs[0].size_with_stride, 0);
341 
342   // Copy input data to input tensor memory
343   int width  = input_attrs[0].dims[2];
344   int stride = input_attrs[0].w_stride;
345   if (width == stride) {
346     memcpy(input_mems[0]->virt_addr, input_data, width * input_attrs[0].dims[1] * input_attrs[0].dims[3]);
347   } else {
348     int height  = input_attrs[0].dims[1];
349     int channel = input_attrs[0].dims[3];
350     // copy from src to dst with stride
351     uint8_t* src_ptr = input_data;
352     uint8_t* dst_ptr = (uint8_t*)input_mems[0]->virt_addr;
353     // width-channel elements
354     int src_wc_elems = width * channel;
355     int dst_wc_elems = stride * channel;
356     for (int h = 0; h < height; ++h) {
357       memcpy(dst_ptr, src_ptr, src_wc_elems);
358       src_ptr += src_wc_elems;
359       dst_ptr += dst_wc_elems;
360     }
361   }
362 
363   // Create output tensor memory
364   rknn_tensor_mem* output_mems[io_num.n_output];
365   for (uint32_t i = 0; i < io_num.n_output; ++i) {
366     output_mems[i] = rknn_create_mem_from_fd(ctx, output_fds[i], output_virts[i], output_attrs[i].size, 0);
367   }
368 
369   // Set input tensor memory
370   ret = rknn_set_io_mem(ctx, input_mems[0], &input_attrs[0]);
371   if (ret < 0) {
372     printf("rknn_set_io_mem fail! ret=%d\n", ret);
373     return -1;
374   }
375 
376   // Set output tensor memory
377   for (uint32_t i = 0; i < io_num.n_output; ++i) {
378     // set output memory and attribute
379     ret = rknn_set_io_mem(ctx, output_mems[i], &output_attrs[i]);
380     if (ret < 0) {
381       printf("rknn_set_io_mem fail! ret=%d\n", ret);
382       return -1;
383     }
384   }
385 
386   // Run
387   printf("Begin perf ...\n");
388   for (int i = 0; i < loop_count; ++i) {
389     int64_t start_us  = getCurrentTimeUs();
390     ret               = rknn_run(ctx, NULL);
391     int64_t elapse_us = getCurrentTimeUs() - start_us;
392     if (ret < 0) {
393       printf("rknn run error %d\n", ret);
394       return -1;
395     }
396     printf("%4d: Elapse Time = %.2fms, FPS = %.2f\n", i, elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
397   }
398 
399   // Get top 5
400   uint32_t topNum = 5;
401   for (uint32_t i = 0; i < io_num.n_output; i++) {
402     uint32_t MaxClass[topNum];
403     float    fMaxProb[topNum];
404     float*   buffer    = (float*)output_mems[i]->virt_addr;
405     uint32_t sz        = output_attrs[i].n_elems;
406     int      top_count = sz > topNum ? topNum : sz;
407 
408     rknn_GetTopN(buffer, fMaxProb, MaxClass, sz, topNum);
409 
410     printf("---- Top%d ----\n", top_count);
411     for (int j = 0; j < top_count; j++) {
412       printf("%8.6f - %d\n", fMaxProb[j], MaxClass[j]);
413     }
414   }
415 
416   // free mb blk memory
417   RK_MPI_MMZ_Free(weight_mb);
418   RK_MPI_MMZ_Free(internal_mb);
419   RK_MPI_MMZ_Free(input_mb);
420   for (uint32_t i = 0; i < io_num.n_output; ++i) {
421     RK_MPI_MMZ_Free(output_mbs[i]);
422   }
423 
424   // Destroy rknn memory
425   rknn_destroy_mem(ctx, input_mems[0]);
426   for (uint32_t i = 0; i < io_num.n_output; ++i) {
427     rknn_destroy_mem(ctx, output_mems[i]);
428   }
429 
430   // destroy
431   rknn_destroy(ctx);
432 
433   free(input_data);
434 
435   return 0;
436 }
437